diff --git a/crates/astra-server-types/src/lib.rs b/crates/astra-server-types/src/lib.rs index 20e01289a5..1acf950f0b 100644 --- a/crates/astra-server-types/src/lib.rs +++ b/crates/astra-server-types/src/lib.rs @@ -2202,12 +2202,14 @@ impl From for AdminUserRoleResponse { #[cfg(feature = "server")] impl From for SessionResponse { fn from(value: SessionRecord) -> Self { + let mut metadata = value.metadata; + metadata.remove(astra_services::auth::session::PROVIDER_SCOPE_PROOF_KEY); Self { session_id: value.session_id, user_id: value.user_id, agent_id: value.agent_id, title: value.title, - metadata: value.metadata, + metadata, status: value.status, event_count: value.event_count, created_at: value.created_at, diff --git a/crates/astra-turn-types/src/lib.rs b/crates/astra-turn-types/src/lib.rs index d74bcabea5..966fd71230 100644 --- a/crates/astra-turn-types/src/lib.rs +++ b/crates/astra-turn-types/src/lib.rs @@ -227,8 +227,9 @@ pub use session_cursor::{ }; pub use session_fork::{ ForkBasisDimensionV1, ForkDimensionDispositionV1, ForkDimensionEvidenceV1, - ForkExcludedAuthorityV1, SESSION_FORK_MANIFEST_SCHEMA_VERSION, SessionForkActivationV1, - SessionForkManifestV1, SessionForkStateV1, SessionForkValidationError, SharedManifestPrefixV1, + ForkExcludedAuthorityV1, ProviderRunOwner, SESSION_FORK_MANIFEST_SCHEMA_VERSION, + SessionForkActivationV1, SessionForkManifestV1, SessionForkStateV1, SessionForkValidationError, + SharedManifestPrefixV1, }; pub use session_handoff::{ HandoffOperationWatermarksV1, HandoffRiskEvidenceV1, MANIFEST_DELTA_SCHEMA_VERSION, diff --git a/crates/astra-turn-types/src/session_fork.rs b/crates/astra-turn-types/src/session_fork.rs index 5862fdb907..2ad5ced09f 100644 --- a/crates/astra-turn-types/src/session_fork.rs +++ b/crates/astra-turn-types/src/session_fork.rs @@ -11,6 +11,13 @@ use crate::{ pub const SESSION_FORK_MANIFEST_SCHEMA_VERSION: u32 = 1; +/// Provider and workspace authenticated by the server, never a client claim. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ProviderRunOwner { + pub provider_id: String, + pub provider_scope_id: String, +} + #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] #[serde(rename_all = "snake_case")] pub enum ForkBasisDimensionV1 { @@ -103,6 +110,10 @@ pub struct SessionForkManifestV1 { pub parent_key: SessionKeyV1, pub child_key: SessionKeyV1, pub parent_head: SessionContextHeadV1, + /// Immutable scope proof for this exact retained prefix and child identity. + /// Older manifests remain readable but cannot establish provider ownership. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_owner: Option, pub dimensions: Vec, pub excluded_authority: Vec, pub state: SessionForkStateV1, @@ -129,6 +140,13 @@ impl SessionForkManifestV1 { return Err(SessionForkValidationError::UnsupportedSchema); } validate_identity(&self.fork_id)?; + if let Some(owner) = &self.provider_owner { + for identity in [&owner.provider_id, &owner.provider_scope_id] { + if identity.is_empty() || identity.trim() != identity { + return Err(SessionForkValidationError::InvalidIdentity); + } + } + } self.shared_prefix().validate_for_child(&self.child_key)?; if self.parent_head.schema_version != SESSION_COORDINATION_SCHEMA_VERSION || self.parent_head.key != self.parent_key @@ -308,6 +326,7 @@ mod tests { fork_id: "fork-a".into(), parent_key: parent_key.clone(), child_key: child_key.clone(), + provider_owner: None, parent_head: SessionContextHeadV1 { schema_version: SESSION_COORDINATION_SCHEMA_VERSION, key: parent_key, diff --git a/crates/runtime/src/server/chat_handlers.rs b/crates/runtime/src/server/chat_handlers.rs index b5aa4fbb6c..ce0534a551 100644 --- a/crates/runtime/src/server/chat_handlers.rs +++ b/crates/runtime/src/server/chat_handlers.rs @@ -31,6 +31,7 @@ pub(super) async fn resolve_or_create_chat_session_id( agent_id, session_id_is_trusted, false, + None, ) .await .map(|resolved| resolved.session_id) @@ -49,6 +50,7 @@ pub(super) async fn resolve_or_create_chat_session( agent_id: Option, session_id_is_trusted: bool, provider_authorized: bool, + provider_principal: Option<&AuthPrincipal>, ) -> Result)> { match requested_session_id { Some(session_id) => { @@ -103,6 +105,40 @@ pub(super) async fn resolve_or_create_chat_session( )]) }); + if provider_authorized && user.user_id.starts_with("provider_authorized:") { + let principal = provider_principal.ok_or_else(|| { + error_response_coded( + StatusCode::FORBIDDEN, + "verified provider identity is required", + "provider_session_creation_forbidden", + ) + })?; + let identity = astra_services::ProviderSessionCreationIdentity::from_principal( + principal, + &format!("chat-{}", uuid::Uuid::new_v4()), + )?; + let result = state + .session_service + .create_provider_session( + identity, + SessionCreateRequestData { + agent_id, + title: None, + metadata, + }, + astra_services::resource_governor::LimitCheck::Allowed, + ) + .await?; + return Ok(ResolvedChatSession { + session_id: Some(result.session.session_id), + full_llm_capture: + crate::turn::llm::exchange_capture::session_full_llm_capture_enabled(Some( + &result.session.metadata, + )), + session_admission_facts: Some(result.admission_facts), + }); + } + match super::session::session_quota::create_session_with_resource_quota( state, user.user_id.clone(), @@ -206,6 +242,7 @@ pub(super) async fn chat_handler( &principal.origin, astra_services::AuthPrincipalOrigin::ProviderAuthorizedRequest(_) ), + Some(&principal), ) .await?; chat_data.session_id = resolved.session_id; @@ -284,6 +321,7 @@ pub(super) async fn chat_stream_handler( &principal.origin, astra_services::AuthPrincipalOrigin::ProviderAuthorizedRequest(_) ), + Some(&principal), ), inject_effective_runtime_context(&state, &principal, &mut chat_data), ); @@ -594,7 +632,8 @@ mod session_resolution_tests { })); state.resource_governor = governor.clone(); let result = - resolve_or_create_chat_session(&state, &test_user(), None, None, false, false).await; + resolve_or_create_chat_session(&state, &test_user(), None, None, false, false, None) + .await; assert!(result.is_err()); assert_eq!(result.err().unwrap().0, StatusCode::SERVICE_UNAVAILABLE); assert_eq!(governor.get_usage("u1").await.sessions_created, 0); @@ -613,6 +652,7 @@ mod session_resolution_tests { Some("agent-1".to_string()), false, false, + None, ) .await .expect("session resolution should succeed"); @@ -674,6 +714,7 @@ mod session_resolution_tests { None, false, false, + None, ) .await .expect("custom session service should resolve an existing session"); @@ -698,6 +739,7 @@ mod session_resolution_tests { None, false, true, + None, ) .await { diff --git a/crates/runtime/src/server/edge/edge_callback_handlers.rs b/crates/runtime/src/server/edge/edge_callback_handlers.rs index c0adf3af81..d6c974d404 100644 --- a/crates/runtime/src/server/edge/edge_callback_handlers.rs +++ b/crates/runtime/src/server/edge/edge_callback_handlers.rs @@ -1504,6 +1504,7 @@ pub(crate) async fn post_provider_interaction_respond_handler( )); } + super::super::run::handlers::authorize_provider_run_session(&state, &principal, run_id).await?; let target = state .execution .run_lifecycle_service @@ -2259,6 +2260,27 @@ mod edge_callback_insert_tests { #[async_trait] impl RunLifecycleService for ApprovalTargetRunLifecycle { + async fn authorize_provider_run_scope( + &self, + user_id: String, + run_id: String, + owner: astra_services::runs::ProviderRunOwner, + ) -> Result<(), (StatusCode, Json)> { + if user_id == "u-approval" + && run_id == self.run_id + && owner.provider_id == "moi" + && owner.provider_scope_id == "workspace-a" + { + Ok(()) + } else { + Err(astra_core::error_response_coded( + StatusCode::FORBIDDEN, + "Session is outside the authorized scope", + "session_fork_forbidden", + )) + } + } + async fn create_run( &self, _user_id: String, diff --git a/crates/runtime/src/server/provider_test_support.rs b/crates/runtime/src/server/provider_test_support.rs index 46a2f9152d..c587f9b154 100644 --- a/crates/runtime/src/server/provider_test_support.rs +++ b/crates/runtime/src/server/provider_test_support.rs @@ -14,6 +14,53 @@ use axum::{Json, Router}; use serde_json::{Value, json}; use tokio::sync::Notify; +/// Exercise provider/binding execution without a persistent Session. This +/// test-only boundary uses production request conversion, runtime injection, +/// lifecycle execution and SSE projection; it does not test inbound HTTP +/// authentication or Session creation/scope admission. +pub async fn ephemeral_provider_chat_stream( + state: &super::AppState, + principal: &astra_services::AuthPrincipal, + payload: Value, +) -> Response { + let request: super::ChatRequest = serde_json::from_value(payload).expect("test chat request"); + assert!( + request.session_id.is_none(), + "ephemeral fixture cannot address a persistent Session" + ); + let mut request = super::chat_request_into_data(request); + let result = async { + super::provider_runtime_context::inject_effective_runtime_context( + state, + principal, + &mut request, + ) + .await?; + state + .execution + .run_lifecycle_service + .stream_chat(principal.user.user_id.clone(), request) + .await + } + .await; + match result { + Ok(mut stream) => super::sse_streaming_response( + stream.session_id, + stream.run_id, + None, + stream + .event_rx + .take() + .expect("a new ephemeral run streams its events"), + ), + Err((status, error)) => super::sse_error_response_from_error_with_context( + status, + error.0, + super::SseErrorContext::default(), + ), + } +} + #[derive(Clone, Debug)] pub struct ProviderRequest { pub method: Method, diff --git a/crates/runtime/src/server/router_builder/sessions.rs b/crates/runtime/src/server/router_builder/sessions.rs index eee47283a2..e7eed92b25 100644 --- a/crates/runtime/src/server/router_builder/sessions.rs +++ b/crates/runtime/src/server/router_builder/sessions.rs @@ -105,6 +105,14 @@ pub(super) fn add_routes(router: Router, state: AppState) -> Router Result<(), (StatusCode, Json)> { + let AuthPrincipalOrigin::ProviderAuthorizedRequest(context) = &principal.origin else { + return Ok(()); + }; + let owner = astra_services::runs::ProviderRunOwner { + provider_id: context.provider_id.clone(), + provider_scope_id: context.provider_scope_id.clone(), + }; + state + .execution + .run_lifecycle_service + .authorize_provider_run_scope(principal.user.user_id.clone(), run_id.to_owned(), owner) + .await +} + pub(crate) async fn stream_run_handler( State(state): State, trace: Option>, @@ -464,6 +483,19 @@ pub(crate) async fn stream_run_handler( } }; + if let Err((status, error)) = authorize_provider_run_session(&state, &principal, &run_id).await + { + return sse_error_response_from_error_with_context( + status, + error.0, + SseErrorContext { + request_id: request_id.as_deref(), + run_id: Some(&run_id), + ..SseErrorContext::default() + }, + ); + } + if query.replay_only { return match state .execution @@ -548,6 +580,7 @@ pub(crate) async fn delete_run_handler( external_request_descriptor(&method, &uri, &headers, "/chat/runs/{run_id}", &body), ) .await?; + authorize_provider_run_session(&state, &principal, &run_id).await?; cancel_run_for_user(&state, run_id, principal.user.user_id).await } diff --git a/crates/runtime/src/server/run/lifecycle/mod.rs b/crates/runtime/src/server/run/lifecycle/mod.rs index 2016d30f41..77f26c6f10 100644 --- a/crates/runtime/src/server/run/lifecycle/mod.rs +++ b/crates/runtime/src/server/run/lifecycle/mod.rs @@ -15175,6 +15175,56 @@ impl AgenticRunLifecycleService { Ok(bindings) } + async fn authorize_fork_session_request( + &self, + user_id: &str, + request: &ChatRequestData, + ) -> Result<(), (StatusCode, Json)> { + if let Some(session_id) = &request.session_id { + self.authorize_fork_session_id( + user_id, + session_id, + request.provider_run_owner.as_ref(), + ) + .await?; + } + Ok(()) + } + + async fn authorize_fork_session_id( + &self, + user_id: &str, + session_id: &str, + provider_owner: Option<&astra_services::runs::ProviderRunOwner>, + ) -> Result<(), (StatusCode, Json)> { + if let Some(pool) = &self.shared_pool { + astra_services::session_fork_coordinator::authorize_provider_fork_session( + pool.get(), + user_id, + session_id, + provider_owner, + Some(self.encryptor.as_ref()), + astra_services::session_fork_coordinator::ProviderSessionScopeRequirement::ExistingConversation, + ) + .await + .map_err(|error| match error { + astra_services::SessionForkCoordinatorError::Forbidden => error_response_coded( + StatusCode::FORBIDDEN, + "Session fork is outside the authorized scope", + "session_fork_forbidden", + ), + error => error_response(StatusCode::INTERNAL_SERVER_ERROR, error.to_string()), + })?; + } else if provider_owner.is_some() { + return Err(error_response_coded( + StatusCode::SERVICE_UNAVAILABLE, + "Provider session scope cannot be verified", + "provider_session_scope_unavailable", + )); + } + Ok(()) + } + async fn prepare_new_owned_execution( &self, user_id: String, @@ -15183,6 +15233,10 @@ impl AgenticRunLifecycleService { streaming: bool, ) -> Result)> { self.require_invocation_composition()?; + if !streaming { + self.authorize_fork_session_request(&user_id, &request) + .await?; + } let request = self.prepare_chat_request(&user_id, request).await?; let request_constraints = self .validate_request_constraints(&user_id, &request) @@ -15325,9 +15379,16 @@ impl AgenticRunLifecycleService { let execution_start_authority = match start_claim { DurableRunStartClaim::Started { authority } => authority, DurableRunStartClaim::Existing { + session_id: existing_session_id, start_request_fingerprint, .. } => { + self.authorize_fork_session_id( + &user_id, + &existing_session_id, + request.provider_run_owner.as_ref(), + ) + .await?; Self::validate_start_request_fingerprint( start_identity .as_ref() @@ -16144,6 +16205,40 @@ impl AgenticRunLifecycleService { #[async_trait] impl RunLifecycleService for AgenticRunLifecycleService { + async fn authorize_provider_run_scope( + &self, + user_id: String, + run_id: String, + owner: astra_services::runs::ProviderRunOwner, + ) -> Result<(), (StatusCode, Json)> { + let pool = self.shared_pool.as_ref().ok_or_else(|| { + error_response_coded( + StatusCode::SERVICE_UNAVAILABLE, + "Provider session scope cannot be verified", + "provider_session_scope_unavailable", + ) + })?; + let session_id: Option = sqlx::query_scalar( + "SELECT session_id FROM agent_runs WHERE user_id = ? AND run_id = ? LIMIT 1", + ) + .bind(&user_id) + .bind(&run_id) + .fetch_optional(pool.get()) + .await + .map_err(|_| { + error_response_coded( + StatusCode::SERVICE_UNAVAILABLE, + "Provider session scope cannot be verified", + "provider_session_scope_unavailable", + ) + })?; + let session_id = session_id.ok_or_else(|| { + error_response_coded(StatusCode::NOT_FOUND, "Run not found", "run_not_found") + })?; + self.authorize_fork_session_id(&user_id, &session_id, Some(&owner)) + .await + } + fn workspace_executor_id(&self) -> Option<&str> { self.selected_server_workspace_provider .as_ref() @@ -16201,6 +16296,8 @@ impl RunLifecycleService for AgenticRunLifecycleService { request: ChatRequestData, ) -> Result)> { self.require_invocation_composition()?; + self.authorize_fork_session_request(&user_id, &request) + .await?; let start_identity = self.run_start_idempotency(&user_id, &request)?; // Replays consume frozen execution facts before mutable admission. if let Some(identity) = start_identity.as_ref() @@ -16208,6 +16305,12 @@ impl RunLifecycleService for AgenticRunLifecycleService { .load_durable_run_for_user(identity.run_id(), &user_id) .await? { + self.authorize_fork_session_id( + &user_id, + &durable.session_id, + request.provider_run_owner.as_ref(), + ) + .await?; Self::validate_start_request_session( identity, request.session_id.as_deref(), diff --git a/crates/runtime/src/server/session/session_handlers.rs b/crates/runtime/src/server/session/session_handlers.rs index 2c7f0f92ee..b607521e51 100644 --- a/crates/runtime/src/server/session/session_handlers.rs +++ b/crates/runtime/src/server/session/session_handlers.rs @@ -441,20 +441,14 @@ pub(crate) async fn create_session_handler( let identity = astra_services::ProviderSessionCreationIdentity::from_principal( &principal, client_ref, )?; - let quota = state - .resource_governor - .check_session_create(&principal.user.user_id) - .await; let result = state .session_service - .create_provider_session(identity, data, quota) + .create_provider_session( + identity, + data, + astra_services::resource_governor::LimitCheck::Allowed, + ) .await?; - if result.created { - state - .resource_governor - .record_session_created(&principal.user.user_id) - .await; - } let status = if result.created { StatusCode::CREATED } else { @@ -2406,18 +2400,155 @@ pub(crate) async fn activate_session_handoff_handler( .map_err(session_handoff_http_error) } +fn fork_provider_scope( + auth: &astra_services::RequestAuthResult, + allow_legacy: bool, +) -> Result, (StatusCode, Json)> { + match &auth.principal.origin { + astra_services::AuthPrincipalOrigin::ProviderAuthorizedRequest(context) => { + match auth.mechanism { + astra_services::RequestAuthMechanism::RequestBoundHmac => { + Ok(Some(astra_services::runs::ProviderRunOwner { + provider_id: context.provider_id.clone(), + provider_scope_id: context.provider_scope_id.clone(), + })) + } + astra_services::RequestAuthMechanism::KnownLegacy if allow_legacy => Ok(None), + _ => Err(session_fork_http_error( + astra_services::SessionForkCoordinatorError::Forbidden, + )), + } + } + _ => Ok(None), + } +} + +pub(crate) async fn get_session_run_cursor_handler( + State(state): State, + Path((session_id, run_id)): Path<(String, String)>, + method: Method, + uri: Uri, + headers: HeaderMap, + body: Bytes, +) -> Result, (StatusCode, Json)> { + let auth = state + .auth_service + .current_principal_for_request_with_provenance( + &headers, + external_request_descriptor( + &method, + &uri, + &headers, + "/sessions/{session_id}/runs/{run_id}/cursor", + &body, + ), + ) + .await?; + let principal = &auth.principal; + state + .session_service + .get_session(session_id.clone(), principal.user.user_id.clone()) + .await?; + let key = server_session_key(&principal.user.user_id, &session_id); + let provider_scope = fork_provider_scope(&auth, false)?; + let cursor = fork_coordinator(&state)? + .resolve_run_cursor(&key, &run_id, provider_scope.as_ref()) + .await + .map_err(session_fork_http_error)?; + Ok(Json(cursor)) +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +pub(crate) struct SessionForkContinuationRequest {} + +pub(crate) async fn activate_session_fork_for_continuation_handler( + State(state): State, + Path((session_id, fork_id)): Path<(String, String)>, + method: Method, + uri: Uri, + headers: HeaderMap, + body: Bytes, +) -> Result, (StatusCode, Json)> { + let auth = state + .auth_service + .current_principal_for_request_with_provenance( + &headers, + external_request_descriptor( + &method, + &uri, + &headers, + "/sessions/{session_id}/forks/{fork_id}/activate-for-continuation", + &body, + ), + ) + .await?; + let provider_owner = fork_provider_scope(&auth, false)?; + let principal = &auth.principal; + let _: SessionForkContinuationRequest = serde_json::from_slice(&body).map_err(|_| { + error_response_coded( + StatusCode::BAD_REQUEST, + "Invalid session fork request", + "session_fork_invalid", + ) + })?; + state + .session_service + .get_session(session_id.clone(), principal.user.user_id.clone()) + .await?; + let key = server_session_key(&principal.user.user_id, &session_id); + let coordinator = fork_coordinator(&state)?; + let manifest = coordinator + .load(&key, &fork_id) + .await + .map_err(session_fork_http_error)?; + if let Some(owner) = provider_owner { + astra_services::session_fork_coordinator::authorize_manifest_provider(&manifest, &owner) + .map_err(session_fork_http_error)?; + } + coordinator + .activate_for_continuation(&key, &fork_id) + .await + .map(Json) + .map_err(session_fork_http_error) +} + pub(crate) async fn prepare_session_fork_handler( State(state): State, Path(session_id): Path, + method: Method, + uri: Uri, headers: HeaderMap, - Json(request): Json, + body: Bytes, ) -> Result, (StatusCode, Json)> { - let user = state.auth_service.current_user(&headers).await?; + let auth = state + .auth_service + .current_principal_for_request_with_provenance( + &headers, + external_request_descriptor( + &method, + &uri, + &headers, + "/sessions/{session_id}/forks", + &body, + ), + ) + .await?; + let principal = &auth.principal; + let request: SessionForkPrepareRequest = serde_json::from_slice(&body).map_err(|_| { + error_response_coded( + StatusCode::BAD_REQUEST, + "Invalid session fork request", + "session_fork_invalid", + ) + })?; + let user = &principal.user; let session = state .session_service .get_session(session_id, user.user_id.clone()) .await?; let parent_key = server_session_key(&user.user_id, &session.session_id); + let provider_owner = fork_provider_scope(&auth, true)?; let child_session_id = deterministic_fork_session_id(&user.user_id, &session.session_id, &request.idempotency_key); let child_key = server_session_key(&user.user_id, &child_session_id); @@ -2448,14 +2579,18 @@ pub(crate) async fn prepare_session_fork_handler( }) .collect(); fork_coordinator(&state)? - .prepare(&astra_services::PrepareSessionForkV1 { - idempotency_key: request.idempotency_key, - parent_key, - child_key, - expected_parent_cursor: request.expected_parent_cursor, - dimensions, - reason: request.reason, - }) + .prepare_admitted( + &astra_services::PrepareSessionForkV1 { + idempotency_key: request.idempotency_key, + parent_key, + child_key, + expected_parent_cursor: request.expected_parent_cursor, + dimensions, + reason: request.reason, + }, + provider_owner.as_ref(), + true, + ) .await .map(Json) .map_err(session_fork_http_error) @@ -2590,17 +2725,49 @@ pub(crate) async fn activate_session_fork_handler( pub(crate) async fn abort_session_fork_handler( State(state): State, Path((session_id, fork_id)): Path<(String, String)>, + method: Method, + uri: Uri, headers: HeaderMap, - Json(request): Json, + body: Bytes, ) -> Result, (StatusCode, Json)> { - let user = state.auth_service.current_user(&headers).await?; + let auth = state + .auth_service + .current_principal_for_request_with_provenance( + &headers, + external_request_descriptor( + &method, + &uri, + &headers, + "/sessions/{session_id}/forks/{fork_id}/abort", + &body, + ), + ) + .await?; + let request: SessionForkAbortRequest = serde_json::from_slice(&body).map_err(|_| { + error_response_coded( + StatusCode::BAD_REQUEST, + "Invalid session fork request", + "session_fork_invalid", + ) + })?; + let user = &auth.principal.user; let session = state .session_service .get_session(session_id, user.user_id.clone()) .await?; - fork_coordinator(&state)? + let key = server_session_key(&user.user_id, &session.session_id); + let coordinator = fork_coordinator(&state)?; + if let Some(owner) = fork_provider_scope(&auth, true)? { + let manifest = coordinator + .load(&key, &fork_id) + .await + .map_err(session_fork_http_error)?; + astra_services::session_fork_coordinator::authorize_manifest_provider(&manifest, &owner) + .map_err(session_fork_http_error)?; + } + coordinator .abort( - &server_session_key(&user.user_id, &session.session_id), + &key, &fork_id, std::time::Duration::from_secs(request.grace_seconds), &request.detail, @@ -2752,10 +2919,13 @@ fn handoff_service( fn fork_coordinator( state: &AppState, ) -> Result<&astra_services::DatabaseSessionForkCoordinator, (StatusCode, Json)> { - state - .session_fork_coordinator - .as_deref() - .ok_or_else(|| internal_error("session fork coordinator is not configured")) + state.session_fork_coordinator.as_deref().ok_or_else(|| { + error_response_coded( + StatusCode::SERVICE_UNAVAILABLE, + "Session fork is temporarily unavailable", + "session_fork_dependency_unavailable", + ) + }) } fn deterministic_fork_session_id( @@ -2852,25 +3022,87 @@ fn session_handoff_http_error( fn session_fork_http_error( error: astra_services::SessionForkCoordinatorError, ) -> (StatusCode, Json) { - match error { - astra_services::SessionForkCoordinatorError::Invalid(message) => { - error_response(StatusCode::BAD_REQUEST, message) - } - astra_services::SessionForkCoordinatorError::NotFound => { - error_response(StatusCode::NOT_FOUND, error.to_string()) - } - astra_services::SessionForkCoordinatorError::Conflict - | astra_services::SessionForkCoordinatorError::IdempotencyMismatch - | astra_services::SessionForkCoordinatorError::WriterConflict => { - error_response(StatusCode::CONFLICT, error.to_string()) + use astra_services::{ + SessionContextCoordinatorError as ContextError, SessionForkCoordinatorError as ForkError, + }; + let (status, code, message) = match &error { + ForkError::Aborted => ( + StatusCode::GONE, + "session_fork_aborted", + "Session fork was cancelled before activation", + ), + ForkError::QuotaDenied { limit, reason } => { + return error_response_coded(StatusCode::TOO_MANY_REQUESTS, reason, limit.error_code()); } - astra_services::SessionForkCoordinatorError::Coordinator(source) => { - session_context_http_error(source) + ForkError::Invalid(_) | ForkError::Coordinator(ContextError::Invalid(_)) => ( + StatusCode::BAD_REQUEST, + "session_fork_invalid", + "Invalid session fork request", + ), + ForkError::Coordinator(ContextError::ResourceLimit(_)) => ( + StatusCode::PAYLOAD_TOO_LARGE, + "session_fork_resource_limit", + "Session fork exceeds the context resource limit", + ), + ForkError::NotFound | ForkError::Coordinator(ContextError::SegmentNotFound) => ( + StatusCode::NOT_FOUND, + "session_fork_not_found", + "Session fork source was not found", + ), + ForkError::Forbidden | ForkError::Coordinator(ContextError::Unauthorized) => ( + StatusCode::FORBIDDEN, + "session_fork_forbidden", + "Session fork is outside the authorized scope", + ), + ForkError::CursorNotReady => ( + StatusCode::CONFLICT, + "session_fork_cursor_not_ready", + "The committed run cursor is not ready", + ), + ForkError::CursorUnavailable => ( + StatusCode::CONFLICT, + "session_fork_cursor_unavailable", + "The committed run cursor is unavailable", + ), + ForkError::CursorUnsupported => ( + StatusCode::UNPROCESSABLE_ENTITY, + "session_fork_cursor_unsupported", + "This run does not support a canonical fork cursor", + ), + ForkError::IdempotencyMismatch + | ForkError::Coordinator(ContextError::IdempotencyMismatch) => ( + StatusCode::CONFLICT, + "session_fork_idempotency_mismatch", + "The fork idempotency key was reused for a different request", + ), + ForkError::WriterConflict => ( + StatusCode::CONFLICT, + "session_fork_writer_conflict", + "The fork writer is held by another actor", + ), + ForkError::Conflict + | ForkError::Coordinator( + ContextError::Fenced + | ContextError::ExecutionOccupied(_) + | ContextError::SessionLifecycleFenced + | ContextError::Expired + | ContextError::DivergentManifest, + ) => ( + StatusCode::CONFLICT, + "session_fork_conflict", + "The fork conflicts with current session state", + ), + _ => { + tracing::warn!(error = %error, "Session fork dependency failed"); + ( + StatusCode::SERVICE_UNAVAILABLE, + "session_fork_dependency_unavailable", + "Session fork is temporarily unavailable", + ) } - _ => internal_error(error), - } + }; + error_response_coded(status, message, code) } - pub(crate) async fn cancel_session_handler( State(state): State, Path(session_id): Path, diff --git a/crates/runtime/src/server/session/session_quota.rs b/crates/runtime/src/server/session/session_quota.rs index 150560bdc0..ac56d0ce61 100644 --- a/crates/runtime/src/server/session/session_quota.rs +++ b/crates/runtime/src/server/session/session_quota.rs @@ -7,7 +7,7 @@ use crate::server::*; /// - session-create quota: counts durable chat sessions (`agent_sessions`) /// - run-start quota: limits execution capacity for active runs /// -/// Keep the session counter here, immediately around the actual session insert. +/// The Session service owns the counter and the actual session insert. /// A single durable session can contain many agentic runs, so run lifecycle code /// must not increment `sessions_created`. pub(crate) async fn create_session_with_resource_quota( @@ -18,17 +18,8 @@ pub(crate) async fn create_session_with_resource_quota( (SessionRecord, Option), (StatusCode, Json), > { - astra_services::resource_governor::enforce_session_create_quota( - &state.resource_governor.check_session_create(&user_id).await, - )?; - - let session = state - .session_service - .create_session_with_admission_facts(user_id.clone(), request) - .await?; state - .resource_governor - .record_session_created(&user_id) - .await; - Ok(session) + .session_service + .create_session_with_resource_quota(user_id, request, state.resource_governor.as_ref()) + .await } diff --git a/crates/runtime/src/server/state_builder/core.rs b/crates/runtime/src/server/state_builder/core.rs index 174a82838d..459b12fa44 100644 --- a/crates/runtime/src/server/state_builder/core.rs +++ b/crates/runtime/src/server/state_builder/core.rs @@ -87,10 +87,13 @@ pub(super) fn build_core_state( shared_pool.clone(), Arc::clone(&session_context_coordinator), )); - let session_fork_coordinator = Arc::new(astra_services::DatabaseSessionForkCoordinator::new( - shared_pool.clone(), - Arc::clone(&session_context_coordinator), - )); + let session_fork_coordinator = Arc::new( + astra_services::DatabaseSessionForkCoordinator::new( + shared_pool.clone(), + Arc::clone(&session_context_coordinator), + ) + .with_scope_encryptor(Arc::clone(shared_encryptor)), + ); AppState::new( ServiceInfo::default(), @@ -109,7 +112,9 @@ pub(super) fn build_core_state( ))) .with_auth_service(auth_service) .with_session_service(Arc::new( - DatabaseSessionService::new(settings.matrixone.clone()).with_pool(shared_pool.clone()), + DatabaseSessionService::new(settings.matrixone.clone()) + .with_pool(shared_pool.clone()) + .with_scope_encryptor(Arc::clone(shared_encryptor)), )) .with_agent_service(Arc::new( DatabaseAgentService::new(settings.matrixone.clone()).with_pool(shared_pool.clone()), diff --git a/crates/runtime/src/server/state_builder/runtime.rs b/crates/runtime/src/server/state_builder/runtime.rs index d403ff8088..312d23e28c 100644 --- a/crates/runtime/src/server/state_builder/runtime.rs +++ b/crates/runtime/src/server/state_builder/runtime.rs @@ -8,9 +8,10 @@ pub(super) async fn build_runtime_wiring( state: &AppState, memoria_port: Option>, ) -> Result> { - let run_store = Arc::new(astra_services::runs::DatabaseRunStateStore::new( - shared_pool.clone(), - )); + let run_store = Arc::new( + astra_services::runs::DatabaseRunStateStore::new(shared_pool.clone()) + .with_scope_encryptor(Arc::clone(run_encryptor)), + ); let run_engine = crate::server::run::engine::RunEngine::new(run_store) .with_metrics_registry(state.metrics_registry()); let selected_workspace_provider = run_engine.execution_owner_pod_id().map(|owner| { diff --git a/crates/runtime/src/server/ws_handler.rs b/crates/runtime/src/server/ws_handler.rs index 6263c0ec4f..819382db9b 100644 --- a/crates/runtime/src/server/ws_handler.rs +++ b/crates/runtime/src/server/ws_handler.rs @@ -740,6 +740,7 @@ async fn handle_chat_message( &conn.principal.origin, astra_services::AuthPrincipalOrigin::ProviderAuthorizedRequest(_) ), + Some(&conn.principal), ) .await { @@ -814,6 +815,9 @@ async fn handle_attach_run( run_id: &str, last_index: u32, ) { + if !authorize_ws_run(socket, state, conn, run_id).await { + return; + } let record = match state .execution .run_lifecycle_service @@ -851,6 +855,22 @@ async fn handle_attach_run( conn.active_run_id = None; } +async fn authorize_ws_run( + socket: &mut WebSocket, + state: &AppState, + conn: &WsConnection, + run_id: &str, +) -> bool { + match super::run::handlers::authorize_provider_run_session(state, &conn.principal, run_id).await + { + Ok(()) => true, + Err((status, error)) => { + send_msg(socket, &ws_error_from_status(status, error.0.detail)).await; + false + } + } +} + async fn inject_ws_effective_runtime_context( state: &AppState, conn: &WsConnection, @@ -866,6 +886,9 @@ async fn handle_cancel_run( conn: &WsConnection, run_id: &str, ) { + if !authorize_ws_run(socket, state, conn, run_id).await { + return; + } match state .execution .run_lifecycle_service @@ -888,6 +911,9 @@ async fn handle_pause_run( run_id: &str, emit_ack: bool, ) { + if !authorize_ws_run(socket, state, conn, run_id).await { + return; + } match state .execution .run_lifecycle_service @@ -918,6 +944,9 @@ async fn handle_resume_run( run_id: &str, emit_ack: bool, ) { + if !authorize_ws_run(socket, state, conn, run_id).await { + return; + } match state .execution .run_lifecycle_service diff --git a/crates/runtime/src/turn/llm/client.rs b/crates/runtime/src/turn/llm/client.rs index 5eabcb8eb7..414ab10160 100644 --- a/crates/runtime/src/turn/llm/client.rs +++ b/crates/runtime/src/turn/llm/client.rs @@ -9470,6 +9470,8 @@ mod tests { #[tokio::test] async fn successful_headers_then_body_budget_records_delivery_unknown() { + // Exercise the request phase after cold client initialization. + let _ = global_llm_client(); let app = Router::new().route( "/chat/completions", post(|| async { @@ -12747,6 +12749,8 @@ mod tests { #[tokio::test] async fn successful_provider_result_survives_slow_terminalization_as_partial_evidence() { + // Exercise the request phase after cold client initialization. + let _ = global_llm_client(); let app = Router::new().route( "/chat/completions", post(|| async { @@ -15567,6 +15571,8 @@ mod tests { #[tokio::test] async fn call_llm_and_collect_cancel_during_429_backoff_sleep() { + // Exercise the request phase after cold client initialization. + let _ = global_llm_client(); reset_rate_limit_cooldown_for_tests(); let hits = Arc::new(AtomicU32::new(0)); let app = Router::new() @@ -15699,6 +15705,8 @@ mod tests { #[tokio::test] async fn call_llm_and_collect_cancel_during_exponential_backoff_after_500() { + // Exercise the request phase after cold client initialization. + let _ = global_llm_client(); reset_rate_limit_cooldown_for_tests(); let hits = Arc::new(AtomicU32::new(0)); let app = Router::new() diff --git a/crates/runtime/tests/moi_fresh_astra.rs b/crates/runtime/tests/moi_fresh_astra.rs index 6378d8957b..69d6cf35d0 100644 --- a/crates/runtime/tests/moi_fresh_astra.rs +++ b/crates/runtime/tests/moi_fresh_astra.rs @@ -27,6 +27,10 @@ use std::{ use tower::ServiceExt; use uuid::Uuid; +#[path = "moi_fresh_astra/fork_auth_compat.rs"] +mod fork_auth_compat; +#[path = "moi_fresh_astra/forks.rs"] +mod forks; #[path = "moi_fresh_astra/normal_chat.rs"] mod normal_chat; #[path = "moi_fresh_astra/reasoning.rs"] diff --git a/crates/runtime/tests/moi_fresh_astra/fork_auth_compat.rs b/crates/runtime/tests/moi_fresh_astra/fork_auth_compat.rs new file mode 100644 index 0000000000..8074a9c27c --- /dev/null +++ b/crates/runtime/tests/moi_fresh_astra/fork_auth_compat.rs @@ -0,0 +1,673 @@ +//! Existing user-token controls remain separate from request-signed MOI forks. +use astra_core::{ + AppSettings, EdgeTokenAuthConfig, ErrorResponse, ProviderRequestAuthConfig, SharedPool, +}; +use astra_runtime::{AppState, build_app, build_server_state}; +use astra_services::{ + AcquireWriterOutcome, AuthLoginRequestData, AuthPrincipal, AuthPrincipalOrigin, + AuthRefreshRequestData, AuthRegisterRequestData, AuthService, AuthTokenRecord, AuthUserRecord, + DatabaseAuthService, DatabaseSessionContextCoordinator, DatabaseSessionForkCoordinator, + DatabaseSessionService, ProviderRequestDescriptor, RequestAuthMechanism, ReserveTurnOutcome, + SessionContextCoordinator, SessionCreateRequestData, SessionService, +}; +use astra_turn_types::{ + ActorContextV1, ActorKindV1, AuthorityEpochsV1, CANONICAL_TURN_DELTA_SCHEMA_VERSION, + CanonicalDeltaModeV1, CanonicalTurnDeltaV1, CoordinatorMutationV1, SessionCursorV1, + SessionKeyV1, SessionSurfaceV1, +}; +use async_trait::async_trait; +use axum::{ + Json, Router, + body::{Body, to_bytes}, + http::{HeaderMap, Request, StatusCode}, + routing::post, +}; +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use hmac::{Hmac, Mac}; +use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; +use sqlx::{Connection, MySqlConnection}; +use std::{ + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; +use tower::ServiceExt; +use uuid::Uuid; + +pub(super) const USER_TOKEN_KEY: &str = "moi-fork-legacy-credential-test-key"; + +pub(super) fn token(claims: Value) -> String { + let claims = URL_SAFE_NO_PAD.encode(claims.to_string()); + let mut signer = Hmac::::new_from_slice(USER_TOKEN_KEY.as_bytes()).unwrap(); + signer.update(claims.as_bytes()); + format!( + "moi-user-token-v1.{claims}.{}", + URL_SAFE_NO_PAD.encode(signer.finalize().into_bytes()) + ) +} + +fn bearer_headers(token: &str) -> HeaderMap { + let mut headers = HeaderMap::new(); + headers.insert("authorization", format!("Bearer {token}").parse().unwrap()); + headers +} + +struct LegacyCase { + name: &'static str, + token: String, + key: SessionKeyV1, + cursor: SessionCursorV1, +} + +impl LegacyCase { + fn prepare_body(&self, intent: &str) -> String { + json!({"idempotency_key":intent,"expected_parent_cursor":self.cursor, + "reason":"existing credential at canonical commit before transcript publication"}) + .to_string() + } + + fn request(&self, path: &str, body: String) -> Request { + Request::builder() + .method("POST") + .uri(path) + .header("content-type", "application/json") + .header("authorization", format!("Bearer {}", self.token)) + .body(Body::from(body)) + .unwrap() + } + + fn prepare_request(&self, intent: &str) -> Request { + self.request( + &format!("/sessions/{}/forks", self.key.session_id), + self.prepare_body(intent), + ) + } +} + +async fn response(app: &Router, request: Request) -> (StatusCode, Value) { + let response = app.clone().oneshot(request).await.unwrap(); + let status = response.status(); + let bytes = to_bytes(response.into_body(), 1024 * 1024).await.unwrap(); + ( + status, + serde_json::from_slice(&bytes).expect("control response JSON"), + ) +} + +struct LegacyFixture { + state: AppState, + app: Router, + auth: Arc, + sessions: Arc, + context: Arc, + forks: Arc, + pool: SharedPool, + database: String, + admin: MySqlConnection, + revocation_task: tokio::task::JoinHandle<()>, + revocation_calls: Arc, + revocation_errors: Arc>>, + cases: Vec, +} + +impl LegacyFixture { + async fn new() -> Self { + assert_eq!(std::env::var("ASTRA_TEST_DB_IT").as_deref(), Ok("1")); + let now = chrono::Utc::now().timestamp(); + let runtime_token = token( + json!({"sub":"fork-runtime-user","workspace_id":"legacy-workspace", + "iss":"moi-core","iat":now,"exp":now+3600,"purpose":"runtime"}), + ); + let edge_token = token( + json!({"sub":"fork-edge-user","workspace_id":"legacy-workspace", + "iss":"moi-backend","iat":now,"exp":now+3600,"purpose":"edge_registration", + "edge_agent_id":"fork-compat-edge","jti":Uuid::new_v4().to_string()}), + ); + let calls = Arc::new(AtomicUsize::new(0)); + let errors = Arc::new(Mutex::new(Vec::new())); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let revocation = Router::new().route( + "/check", + post({ + let calls = calls.clone(); + let errors = errors.clone(); + let expected = edge_token.clone(); + move |Json(body): Json| { + let calls = calls.clone(); + let errors = errors.clone(); + let expected = expected.clone(); + async move { + calls.fetch_add(1, Ordering::SeqCst); + let valid = body == json!({"token":expected}); + if !valid { + errors.lock().unwrap().push( + "revocation request did not carry the real edge credential".into(), + ); + } + Json(json!({"code":if valid {0} else {1},"data":{"ok":valid}})) + } + } + }), + ); + let revocation_task = tokio::spawn(async move { + axum::serve(listener, revocation).await.unwrap(); + }); + let mut settings = AppSettings::from_explicit_env().expect("explicit test settings"); + let database = format!("astra_test_moi_{}", Uuid::new_v4().simple()); + settings.matrixone.database = "mysql".into(); + let mut admin = MySqlConnection::connect(&settings.matrixone.database_url_with_password()) + .await + .unwrap(); + sqlx::query(sqlx::AssertSqlSafe(format!("CREATE DATABASE `{database}`"))) + .execute(&mut admin) + .await + .unwrap(); + eprintln!("MOI legacy fork auth fixture created: {database}"); + settings.matrixone.database = database.clone(); + settings.memoria.base_url = endpoint.clone(); + settings.external_providers.clear(); + settings.edge_token_auth = EdgeTokenAuthConfig { + key: USER_TOKEN_KEY.into(), + check_endpoint: format!("{endpoint}/check"), + }; + let state = build_server_state(settings.clone()) + .await + .expect("fresh real Astra state"); + let pool = state.shared_pool.as_ref().unwrap().clone(); + let auth = Arc::new( + DatabaseAuthService::new(settings.matrixone.clone(), settings.jwt.clone()) + .with_pool(pool.clone()) + .with_edge_token_auth(settings.edge_token_auth.clone()), + ); + let sessions = + Arc::new(DatabaseSessionService::new(settings.matrixone).with_pool(pool.clone())); + let context = Arc::new(DatabaseSessionContextCoordinator::new(pool.clone())); + let forks = Arc::new(DatabaseSessionForkCoordinator::new( + pool.clone(), + context.clone(), + )); + let state = state + .with_auth_service(auth.clone()) + .with_session_service(sessions.clone()) + .with_session_fork_coordinator(forks.clone()); + let app = build_app(state.clone()); + let mut cases = Vec::new(); + for (name, token, edge) in [ + ("runtime", runtime_token, false), + ("edge-registration", edge_token, true), + ] { + let headers = bearer_headers(&token); + let principal = auth + .current_principal(&headers) + .await + .expect("real token verification"); + assert!(matches!( + &principal.origin, + AuthPrincipalOrigin::ProviderAuthorizedRequest(_) + )); + assert_eq!(principal.is_edge_registration(), edge); + assert_eq!(auth.current_user(&headers).await.unwrap(), principal.user); + let request = SessionCreateRequestData { + agent_id: None, + title: Some(format!("credential fork auth {name}")), + metadata: None, + }; + let session = sessions + .create_session(principal.user.user_id.clone(), request) + .await + .expect("canonical session owner creates parent"); + let key = SessionKeyV1::owner_session( + "server", + &principal.user.user_id, + &session.session_id, + "main", + ); + let actor = ActorContextV1::owner_user( + &principal.user.user_id, + format!("legacy-{name}"), + ActorKindV1::Server, + SessionSurfaceV1::Server, + None, + AuthorityEpochsV1::default(), + ); + let lease = match context + .acquire_writer(&key, None, &actor, Duration::from_secs(60), "legacy-writer") + .await + .expect("acquire canonical writer") + { + AcquireWriterOutcome::Acquired(lease) => lease, + other => panic!("unexpected writer result {other:?}"), + }; + let reservation = match context + .reserve_turn(&lease, None, Duration::from_secs(30), "legacy-turn", None) + .await + .expect("reserve canonical turn") + { + ReserveTurnOutcome::Reserved(reservation) => reservation, + other => panic!("unexpected reservation result {other:?}"), + }; + // The canonical commit is valid before the separate Server lifecycle + // publishes agent_runs/transcript evidence. Do not hand-edit any + // marker or principal: only the original canonical owner writes here. + let delta = CanonicalTurnDeltaV1 { + schema_version: CANONICAL_TURN_DELTA_SCHEMA_VERSION, + completed_turn: 1, + journal_event_seq: 1, + conversation_seq: 1, + compaction_generation: 0, + config_version_id: None, + mode: CanonicalDeltaModeV1::Append, + logical_segments: vec![vec![ + json!({"role":"user","content":format!("legacy {name} question")}), + json!({"role":"assistant","content":format!("legacy {name} answer")}), + ]], + }; + let cursor = match context + .commit_turn_and_release_writer(&reservation, delta, "legacy-commit") + .await + .expect("commit canonical head without Server projection") + { + CoordinatorMutationV1::Applied { cursor } => cursor, + other => panic!("unexpected commit result {other:?}"), + }; + cases.push(LegacyCase { + name, + token, + key, + cursor, + }); + } + let fixture = Self { + state, + app, + auth, + sessions, + context, + forks, + pool, + database, + admin, + revocation_task, + revocation_calls: calls, + revocation_errors: errors, + cases, + }; + fixture.assert_canonical_without_server_projection().await; + fixture + } + + async fn assert_canonical_without_server_projection(&self) { + for case in &self.cases { + self.sessions + .get_session(case.key.session_id.clone(), case.key.owner_user_id.clone()) + .await + .expect("parent remains owned and live"); + let head = self + .context + .load_head(&case.key) + .await + .unwrap() + .expect("committed head"); + assert_eq!(head.cursor, case.cursor); + assert!(head.total_canonical_bytes > 0); + } + for table in ["agent_runs", "agent_run_events", "session_transcript_items"] { + let count: i64 = sqlx::query_scalar(sqlx::AssertSqlSafe(format!( + "SELECT COUNT(*) FROM `{table}`" + ))) + .fetch_one(self.pool.get()) + .await + .unwrap(); + assert_eq!( + count, 0, + "fixture must not invent Server lifecycle evidence in {table}" + ); + } + assert!( + self.revocation_calls.load(Ordering::SeqCst) > 0, + "real loopback revocation was exercised" + ); + assert!(self.revocation_errors.lock().unwrap().is_empty()); + } + + async fn assert_no_forks_or_pins(&self) { + for case in &self.cases { + assert!(self.forks.list(&case.key, 10).await.unwrap().is_empty()); + } + let pins: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM conversation_manifest_pins") + .fetch_one(self.pool.get()) + .await + .unwrap(); + assert_eq!(pins, 0, "denied auth cannot pin or allocate a child"); + let sessions: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM agent_sessions") + .fetch_one(self.pool.get()) + .await + .unwrap(); + assert_eq!(sessions, 2, "denied auth cannot create a child Session"); + self.assert_canonical_without_server_projection().await; + } + + async fn cleanup(mut self) { + self.state.close_database_pools().await; + self.revocation_task.abort(); + sqlx::query(sqlx::AssertSqlSafe(format!( + "DROP DATABASE `{}`", + self.database + ))) + .execute(&mut self.admin) + .await + .unwrap(); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "real disposable MatrixOne; existing moi-user-token fork admission"] +async fn legacy_runtime_and_edge_tokens_prepare_before_server_projection() { + let fixture = LegacyFixture::new().await; + for case in &fixture.cases { + let mut forged: Value = + serde_json::from_str(&case.prepare_body("foreign-history-root")).unwrap(); + forged["expected_parent_cursor"]["canonical_root_hash"] = json!("0".repeat(64)); + let path = format!("/sessions/{}/forks", case.key.session_id); + let (status, error) = response(&fixture.app, case.request(&path, forged.to_string())).await; + assert_eq!(status, StatusCode::CONFLICT, "{error}"); + assert_eq!(error["error_code"], "session_fork_conflict"); + } + fixture.assert_no_forks_or_pins().await; + let mut observed = Vec::new(); + for case in &fixture.cases { + let result = response( + &fixture.app, + case.prepare_request("legacy-compatible-intent"), + ) + .await; + eprintln!("legacy {} prepare: {} {}", case.name, result.0, result.1); + observed.push(result); + } + // Record both real mechanisms before asserting; an edge-only exemption + // must not hide the runtime-token compatibility regression. + let mut outcomes = Vec::new(); + for (case, (status, manifest)) in fixture.cases.iter().zip(observed) { + let activation = if status == StatusCode::OK { + let path = format!( + "/sessions/{}/forks/{}/activate", + case.key.session_id, + manifest["fork_id"].as_str().unwrap() + ); + Some( + response( + &fixture.app, + case.request( + &path, + json!({"placement":"server","writer_ttl_seconds":30}).to_string(), + ), + ) + .await, + ) + } else { + None + }; + outcomes.push(( + case.name, + serde_json::to_value(&case.cursor).unwrap(), + status, + manifest, + activation, + )); + } + fixture.assert_canonical_without_server_projection().await; + fixture.cleanup().await; + for (name, cursor, status, manifest, activation) in outcomes { + assert_eq!( + status, + StatusCode::OK, + "{name} previously admitted prepare: {manifest}" + ); + assert_eq!(manifest["parent_head"]["cursor"], cursor); + let (status, active) = activation.expect("successful prepare must be activated"); + assert_eq!(status, StatusCode::OK, "{name} continuation: {active}"); + assert_eq!(active["manifest"]["state"], "active"); + assert_eq!( + active["writer_lease"]["key"], + active["manifest"]["child_key"] + ); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "real disposable MatrixOne; new provider fork controls require request signing"] +async fn legacy_runtime_and_edge_tokens_cannot_use_new_provider_fork_controls() { + let fixture = LegacyFixture::new().await; + let mut observed = Vec::new(); + for case in &fixture.cases { + for (method, path, body) in [ + ( + "GET", + format!( + "/sessions/{}/runs/{}/cursor", + case.key.session_id, + Uuid::new_v4() + ), + "", + ), + ( + "POST", + format!( + "/sessions/{}/forks/{}/activate-for-continuation", + case.key.session_id, + Uuid::new_v4() + ), + "{}", + ), + ] { + let request = Request::builder() + .method(method) + .uri(path) + .header("authorization", format!("Bearer {}", case.token)) + .header("content-type", "application/json") + .body(Body::from(body)) + .unwrap(); + let (status, body) = response(&fixture.app, request).await; + observed.push((case.name, status, body)); + } + } + fixture.assert_no_forks_or_pins().await; + fixture.cleanup().await; + for (name, status, body) in observed { + assert_eq!(status, StatusCode::FORBIDDEN, "{name}: {body}"); + } +} + +struct RequestOnlyProviderAuth { + real: Arc, + request_calls: AtomicUsize, +} + +#[async_trait] +impl AuthService for RequestOnlyProviderAuth { + async fn register( + &self, + request: AuthRegisterRequestData, + ) -> Result)> { + self.real.register(request).await + } + async fn login( + &self, + request: AuthLoginRequestData, + ) -> Result)> { + self.real.login(request).await + } + async fn refresh( + &self, + request: AuthRefreshRequestData, + ) -> Result)> { + self.real.refresh(request).await + } + async fn logout( + &self, + request: AuthRefreshRequestData, + ) -> Result<(), (StatusCode, Json)> { + self.real.logout(request).await + } + async fn current_user( + &self, + _headers: &HeaderMap, + ) -> Result)> { + panic!("request-only verifier must not fall back to current_user") + } + async fn current_principal_for_request( + &self, + headers: &HeaderMap, + request: ProviderRequestDescriptor, + ) -> Result)> { + self.request_calls.fetch_add(1, Ordering::SeqCst); + self.real + .current_principal_for_request(headers, request) + .await + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "real disposable MatrixOne; request-only provider adapter remains conservative"] +async fn request_only_provider_default_cannot_bypass_scope_or_retry_legacy_auth() { + let fixture = LegacyFixture::new().await; + let auth = Arc::new(RequestOnlyProviderAuth { + real: fixture.auth.clone(), + request_calls: AtomicUsize::new(0), + }); + let app = build_app(fixture.state.clone().with_auth_service(auth.clone())); + for case in &fixture.cases { + let (status, error) = response( + &app, + case.prepare_request("unknown-mechanism-must-be-scoped"), + ) + .await; + assert_eq!( + status, + StatusCode::FORBIDDEN, + "unspecified provider mechanism: {error}" + ); + assert_eq!(error["error_code"], "session_fork_forbidden"); + } + assert_eq!( + auth.request_calls.load(Ordering::SeqCst), + 2, + "verify once per request" + ); + fixture.assert_no_forks_or_pins().await; + for case in &fixture.cases { + let mut request = case.prepare_request("malformed-provider-headers"); + request + .headers_mut() + .insert("x-astra-provider", "moi".parse().unwrap()); + let (status, error) = response(&fixture.app, request).await; + assert_eq!( + status, + StatusCode::BAD_REQUEST, + "partial provider headers cannot fall back: {error}" + ); + } + assert_real_request_auth_provenance(&fixture).await; + fixture.assert_no_forks_or_pins().await; + fixture.cleanup().await; +} + +// Exercise the real DB verifier's mechanism evidence and shared durable replay +// guard. The first valid HMAC must succeed; verifying twice would consume its +// nonce before this result could be returned. +async fn assert_real_request_auth_provenance(fixture: &LegacyFixture) { + let body = "{}"; + let descriptor = ProviderRequestDescriptor { + method: "POST".into(), + path: format!("/sessions/{}/forks", fixture.cases[0].key.session_id), + route: Some("/sessions/{session_id}/forks".into()), + request_id: Some(Uuid::new_v4().to_string()), + body_digest: Some(format!("sha256:{:x}", Sha256::digest(body.as_bytes()))), + }; + for case in &fixture.cases { + let result = fixture + .auth + .current_principal_for_request_with_provenance( + &bearer_headers(&case.token), + descriptor.clone(), + ) + .await + .unwrap(); + assert_eq!( + result.mechanism, + RequestAuthMechanism::KnownLegacy, + "{} provenance", + case.name + ); + assert_eq!(result.principal.user.user_id, case.key.owner_user_id); + } + let auth = fixture + .auth + .as_ref() + .clone() + .with_provider_request_auth(vec![ProviderRequestAuthConfig { + provider: "moi".into(), + auth_type: "hmac".into(), + key: USER_TOKEN_KEY.into(), + }]); + let now = chrono::Utc::now().timestamp(); + let claims = json!({"provider":"moi","sub":"fork-hmac-user","scope":"legacy-workspace", + "method":descriptor.method,"path":descriptor.path,"route":descriptor.route, + "request_id":descriptor.request_id,"body_digest":descriptor.body_digest, + "nonce":Uuid::new_v4().to_string(),"iat":now,"exp":now+300}); + let claims = URL_SAFE_NO_PAD.encode(claims.to_string()); + let mut signer = Hmac::::new_from_slice(USER_TOKEN_KEY.as_bytes()).unwrap(); + signer.update(claims.as_bytes()); + let token = format!( + "moi-provider-v1.{claims}.{}", + URL_SAFE_NO_PAD.encode(signer.finalize().into_bytes()) + ); + let mut headers = bearer_headers(&token); + headers.insert("x-astra-provider", "moi".parse().unwrap()); + headers.insert( + "x-astra-provider-action", + "authorize_request".parse().unwrap(), + ); + let result = auth + .current_principal_for_request_with_provenance(&headers, descriptor.clone()) + .await + .expect("first valid HMAC verifies and consumes its nonce once"); + assert_eq!(result.mechanism, RequestAuthMechanism::RequestBoundHmac); + assert!(matches!( + result.principal.origin, + AuthPrincipalOrigin::ProviderAuthorizedRequest(_) + )); + let replay = auth + .current_principal_for_request(&headers, descriptor.clone()) + .await + .unwrap_err(); + assert_eq!( + replay.0, + StatusCode::UNAUTHORIZED, + "old and new methods share one replay guard" + ); + let rows: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM auth_provider_request_replay WHERE request_id = ?", + ) + .bind(descriptor.request_id.as_ref().unwrap()) + .fetch_one(fixture.pool.get()) + .await + .unwrap(); + assert_eq!( + rows, 1, + "exactly one durable authorization survives the replay rejection" + ); + let stripped = auth + .current_principal_for_request_with_provenance(&bearer_headers(&token), descriptor) + .await + .unwrap_err(); + assert_eq!( + stripped.0, + StatusCode::UNAUTHORIZED, + "removing HMAC headers cannot turn it into a legacy credential" + ); +} diff --git a/crates/runtime/tests/moi_fresh_astra/forks.rs b/crates/runtime/tests/moi_fresh_astra/forks.rs new file mode 100644 index 0000000000..77ee2bb531 --- /dev/null +++ b/crates/runtime/tests/moi_fresh_astra/forks.rs @@ -0,0 +1,1633 @@ +//! New Session fork consumer contract. The ordinary MOI chat baseline stays frozen. +use super::*; +use astra_runtime::server::provider_test_support::{ + ProviderGateway, ProviderResponse, ProviderScript, +}; +use astra_services::{ + DatabaseSessionContextCoordinator, DatabaseSessionForkCoordinator, DatabaseSessionService, + SessionContextCoordinator, SessionService, +}; +use astra_turn_types::SessionKeyV1; + +async fn control( + app: &Router, + method: &str, + path: &str, + route: &str, + body: &str, +) -> (StatusCode, Value) { + response(app, signed_http_request(method, path, route, body.into())).await +} + +async fn response(app: &Router, request: Request) -> (StatusCode, Value) { + let response = app.clone().oneshot(request).await.unwrap(); + let status = response.status(); + let body = to_bytes(response.into_body(), 1024 * 1024).await.unwrap(); + ( + status, + serde_json::from_slice(&body).expect("control API JSON"), + ) +} + +fn signed_in_scope( + method: &str, + path: &str, + route: &str, + body: &str, + scope: &str, +) -> Request { + let mut request = signed_http_request(method, path, route, body.into()); + let token = request.headers()["authorization"] + .to_str() + .unwrap() + .strip_prefix("Bearer moi-provider-v1.") + .unwrap(); + let encoded = token.split('.').next().unwrap(); + let mut claims: Value = + serde_json::from_slice(&URL_SAFE_NO_PAD.decode(encoded).unwrap()).unwrap(); + claims["scope"] = json!(scope); + let encoded = URL_SAFE_NO_PAD.encode(claims.to_string()); + let mut signer = Hmac::::new_from_slice(KEY.as_bytes()).unwrap(); + signer.update(encoded.as_bytes()); + let token = format!( + "Bearer moi-provider-v1.{encoded}.{}", + URL_SAFE_NO_PAD.encode(signer.finalize().into_bytes()) + ); + request + .headers_mut() + .insert("authorization", token.parse().unwrap()); + request +} + +async fn cursor(app: &Router, session: &str, run: &str) -> Value { + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let (status, value) = control( + app, + "GET", + &format!("/sessions/{session}/runs/{run}/cursor"), + "/sessions/{session_id}/runs/{run_id}/cursor", + "", + ) + .await; + if value["error_code"] == "session_fork_cursor_not_ready" { + tokio::time::sleep(Duration::from_millis(20)).await; + continue; + } + assert_eq!(status, StatusCode::OK, "cursor: {value}"); + return value; + } + }) + .await + .unwrap_or_else(|_| panic!("committed projection unavailable: session {session}, run {run}")) +} + +async fn fork(app: &Router, parent: &str, cursor: &Value, key: &str) -> (Value, String) { + let path = format!("/sessions/{parent}/forks"); + let payload = json!({"idempotency_key":key,"expected_parent_cursor":cursor,"reason":"MOI selected completed turn"}).to_string(); + let (status, manifest) = + control(app, "POST", &path, "/sessions/{session_id}/forks", &payload).await; + assert_eq!(status, StatusCode::OK, "prepare: {manifest}"); + let (status, replay) = + control(app, "POST", &path, "/sessions/{session_id}/forks", &payload).await; + assert_eq!(status, StatusCode::OK); + assert_eq!( + manifest, replay, + "same intent cannot allocate a second fork" + ); + let id = manifest["fork_id"].as_str().unwrap(); + let activate = format!("/sessions/{parent}/forks/{id}/activate-for-continuation"); + let (status, active) = control( + app, + "POST", + &activate, + "/sessions/{session_id}/forks/{fork_id}/activate-for-continuation", + "{}", + ) + .await; + assert_eq!(status, StatusCode::OK, "activate: {active}"); + assert_eq!(active["manifest"]["state"], "active"); + assert!( + active.get("writer_lease").is_none(), + "activation must not grant execution authority" + ); + let child = manifest["child_key"]["session_id"] + .as_str() + .unwrap() + .to_owned(); + assert_eq!(active["child_head"]["cursor"]["session_id"], child); + (active, child) +} + +fn answer(text: &str) -> ProviderResponse { + ProviderResponse::OpenAi(json!({"id":"fork-answer","model":"moi-contract-model", + "choices":[{"index":0,"message":{"role":"assistant","content":text},"finish_reason":"stop"}], + "usage":{"prompt_tokens":100,"completion_tokens":10,"total_tokens":110}})) +} + +struct ReviewFixture { + state: astra_runtime::AppState, + app: Router, + moi: MoiClient, + provider: ProviderGateway, + sessions: DatabaseSessionService, + admin: MySqlConnection, + database: String, + gateway_task: tokio::task::JoinHandle<()>, + parent: String, + parent_run: String, + parent_cursor: Value, +} + +impl ReviewFixture { + async fn new() -> Self { + assert_eq!(std::env::var("ASTRA_TEST_DB_IT").as_deref(), Ok("1")); + let mut settings = AppSettings::from_explicit_env().unwrap(); + settings.matrixone.database = "mysql".into(); + let mut admin = MySqlConnection::connect(&settings.matrixone.database_url_with_password()) + .await + .unwrap(); + let database = format!("astra_test_moi_fork_{}", Uuid::new_v4().simple()); + eprintln!("MOI review fixture database: {database}"); + sqlx::query(sqlx::AssertSqlSafe(format!("CREATE DATABASE `{database}`"))) + .execute(&mut admin) + .await + .unwrap(); + settings.matrixone.database = database.clone(); + settings.provider_request_auth = vec![ProviderRequestAuthConfig { + provider: "moi".into(), + auth_type: "hmac".into(), + key: KEY.into(), + }]; + settings.external_providers.clear(); + let gateway_state = Arc::new(Mutex::new(GatewayState::default())); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + settings.memoria.base_url = endpoint.clone(); + let gateway_app = gateway_router(gateway_state.clone()); + let gateway_task = tokio::spawn(async move { + axum::serve(listener, gateway_app).await.unwrap(); + }); + let session_settings = settings.matrixone.clone(); + let state = build_server_state(settings).await.unwrap(); + let sessions = DatabaseSessionService::new(session_settings) + .with_pool(state.shared_pool.as_ref().unwrap().clone()); + let app = build_app(state.clone()); + let mut moi = MoiClient::new(false, endpoint, "fork-review", gateway_state); + moi.repair_session(&app).await; + moi.repair_bindings(&app).await; + moi.history.clear(); + let provider = ProviderGateway::start(vec![ProviderScript::new( + "fork review", + |r| r.path == "/v1/chat/completions", + (0..5).map(|_| answer("REVIEW-ANSWER")).collect(), + )]) + .await; + moi.model_endpoint = Some(format!("{}/v1/chat/completions", provider.base_url)); + let parent = moi.session.clone().unwrap(); + let run = moi + .chat(&app, "WORKSPACE-A-PRIVATE-HISTORY", "REVIEW-ANSWER") + .await; + let parent_cursor = cursor(&app, &parent, &run).await; + Self { + state, + app, + moi, + provider, + sessions, + admin, + database, + gateway_task, + parent, + parent_run: run, + parent_cursor, + } + } + + async fn close(mut self) { + assert!( + self.state + .drain_background_runs(Duration::from_secs(10)) + .await + ); + self.state.close_database_pools().await; + self.gateway_task.abort(); + sqlx::query(sqlx::AssertSqlSafe(format!( + "DROP DATABASE `{}`", + self.database + ))) + .execute(&mut self.admin) + .await + .unwrap(); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "real disposable MatrixOne; provider run stream and cancel scope"] +async fn provider_fork_review_rejects_cross_workspace_run_stream_and_cancel() { + let fixture = ReviewFixture::new().await; + let run = fixture.parent_run.clone(); + let stream_path = format!("/chat/runs/{run}/stream"); + let replay_path = format!("{stream_path}?replay_only=true"); + let mut rejected_streams = Vec::new(); + for path in [&stream_path, &replay_path] { + let mut request = signed_in_scope( + "GET", + &stream_path, + "/chat/runs/{run_id}/stream", + "", + "other-workspace", + ); + *request.uri_mut() = path.parse().unwrap(); + let response = fixture.app.clone().oneshot(request).await.unwrap(); + let status = response.status(); + let bytes = to_bytes(response.into_body(), 1024 * 1024).await.unwrap(); + let body = String::from_utf8(bytes.to_vec()).unwrap(); + rejected_streams.push((status, body)); + } + let cancel_path = format!("/chat/runs/{run}"); + let (cancel_status, cancel_body) = self::response( + &fixture.app, + signed_in_scope( + "DELETE", + &cancel_path, + "/chat/runs/{run_id}", + "", + "other-workspace", + ), + ) + .await; + // The same-workspace replay must remain usable after a rejected cancellation. + let mut request = + signed_http_request("GET", &stream_path, "/chat/runs/{run_id}/stream", "".into()); + *request.uri_mut() = replay_path.parse().unwrap(); + let response = fixture.app.clone().oneshot(request).await.unwrap(); + let status = response.status(); + let body = to_bytes(response.into_body(), 1024 * 1024).await.unwrap(); + let body = String::from_utf8(body.to_vec()).unwrap(); + fixture.close().await; + for (status, body) in rejected_streams { + assert!( + status == StatusCode::FORBIDDEN + || events(&body) + .iter() + .any(|event| event["error_code"] == "session_fork_forbidden"), + "{status} {body}" + ); + assert!(!body.contains("WORKSPACE-A-PRIVATE-HISTORY")); + } + assert_eq!(cancel_status, StatusCode::FORBIDDEN, "{cancel_body}"); + assert_eq!(status, StatusCode::OK, "{body}"); + assert!( + events(&body) + .iter() + .any(|event| event["type"] == "run_finished"), + "{body}" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "real disposable MatrixOne; PR review workspace regression"] +async fn provider_fork_review_rejects_cross_workspace_chat() { + let mut fixture = ReviewFixture::new().await; + let (active, child) = fork( + &fixture.app, + &fixture.parent, + &fixture.parent_cursor, + "scope-child", + ) + .await; + fixture.moi.session = Some(child.clone()); + let owner = fixture.parent_cursor["owner_id"] + .as_str() + .unwrap() + .to_owned(); + fixture + .sessions + .update_session( + child.clone(), + owner.clone(), + astra_services::SessionUpdateRequestData { + title: None, + status: None, + metadata_patch: None, + metadata: Some(serde_json::Map::from_iter([( + "provider_owner".into(), + json!({"provider_id":"moi", "provider_scope_id":"other-workspace"}), + )])), + }, + ) + .await + .unwrap(); + let body = fixture.moi.payload("WORKSPACE-B-REQUEST").to_string(); + let response = fixture + .app + .clone() + .oneshot(signed_in_scope( + "POST", + "/chat/stream", + "/chat/stream", + &body, + "other-workspace", + )) + .await + .unwrap(); + let status = response.status(); + let body = to_bytes(response.into_body(), 1024 * 1024).await.unwrap(); + let body = String::from_utf8(body.to_vec()).unwrap(); + let requests = fixture.provider.requests.lock().await.len(); + let (background_status, background_error) = self::response( + &fixture.app, + signed_in_scope( + "POST", + "/chat", + "/chat", + &fixture.moi.payload("B-BACKGROUND").to_string(), + "other-workspace", + ), + ) + .await; + assert_eq!( + background_status, + StatusCode::FORBIDDEN, + "{background_error}" + ); + let mut replay_request = fixture.moi.payload("SAME-WORKSPACE-CONTROL"); + replay_request["context"] = json!({"task_ref":"fork-scope-replay"}); + let (same_status, same_body) = post(&fixture.app, "/chat/stream", &replay_request).await; + assert_eq!(same_status, StatusCode::OK); + assert!( + events(&same_body) + .iter() + .any(|event| event["type"] == "run_finished"), + "{same_body}" + ); + // Model a pre-upgrade provider fork: even exact Run replay must not bypass + // its missing immutable scope proof. + let mut old_manifest = active["manifest"].clone(); + old_manifest + .as_object_mut() + .unwrap() + .remove("provider_owner"); + sqlx::query( + "UPDATE session_forks SET manifest_json = ? WHERE owner_user_id = ? AND fork_id = ?", + ) + .bind(old_manifest.to_string()) + .bind(&owner) + .bind(old_manifest["fork_id"].as_str().unwrap()) + .execute(fixture.state.shared_pool.as_ref().unwrap().get()) + .await + .unwrap(); + let (old_status, old_body) = post(&fixture.app, "/chat/stream", &replay_request).await; + assert!( + old_status == StatusCode::FORBIDDEN + || events(&old_body) + .iter() + .any(|event| event["error_code"] == "session_fork_forbidden"), + "{old_body}" + ); + assert_eq!(fixture.provider.requests.lock().await.len(), 2); + fixture.close().await; + assert!( + status == StatusCode::FORBIDDEN + || events(&body) + .iter() + .any(|event| event["error_code"] == "session_fork_forbidden"), + "cross-workspace continuation: {status} {body}" + ); + assert_eq!( + requests, 1, + "rejected continuation must never send inherited history to a model" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "real disposable MatrixOne; PR review retained-head regression"] +async fn provider_fork_review_reforks_inherited_head_after_middle_deletion() { + let mut fixture = ReviewFixture::new().await; + let (_, middle_id) = fork( + &fixture.app, + &fixture.parent, + &fixture.parent_cursor, + "middle", + ) + .await; + fixture.moi.session = Some(middle_id.clone()); + let middle_run = fixture + .moi + .chat(&fixture.app, "MIDDLE-HISTORY", "REVIEW-ANSWER") + .await; + let middle_cursor = cursor(&fixture.app, &middle_id, &middle_run).await; + let (child, child_id) = fork(&fixture.app, &middle_id, &middle_cursor, "child").await; + let (_, control_id) = fork( + &fixture.app, + &child_id, + &child["child_head"]["cursor"], + "before-delete", + ) + .await; + let owner = fixture.parent_cursor["owner_id"].as_str().unwrap(); + fixture + .sessions + .delete_session(control_id, owner.into()) + .await + .unwrap(); + fixture + .sessions + .delete_session(middle_id, owner.into()) + .await + .unwrap(); + let payload = json!({"idempotency_key":"after-delete", "expected_parent_cursor":child["child_head"]["cursor"], "reason":"retained exact head"}).to_string(); + let (status, value) = control( + &fixture.app, + "POST", + &format!("/sessions/{child_id}/forks"), + "/sessions/{session_id}/forks", + &payload, + ) + .await; + fixture.close().await; + assert_eq!( + status, + StatusCode::OK, + "retained inherited head must remain authorized: {value}" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "real disposable MatrixOne; ordinary parent scope laundering regression"] +async fn provider_fork_review_rejects_history_laundering_through_ordinary_parent() { + let fixture = ReviewFixture::new().await; + let response = fixture + .app + .clone() + .oneshot(signed_in_scope( + "POST", + "/chat/stream", + "/chat/stream", + &fixture.moi.payload("B-PARENT-CONTINUATION").to_string(), + "other-workspace", + )) + .await + .unwrap(); + let status = response.status(); + let bytes = to_bytes(response.into_body(), 1024 * 1024).await.unwrap(); + let body = String::from_utf8(bytes.to_vec()).unwrap(); + let parsed = events(&body); + let mut fork_status = None; + if let Some(run) = parsed + .iter() + .find(|event| event["type"] == "session_info") + .and_then(|event| event["run_id"].as_str()) + { + let path = format!("/sessions/{}/runs/{run}/cursor", fixture.parent); + let (cursor_status, cursor) = self::response( + &fixture.app, + signed_in_scope( + "GET", + &path, + "/sessions/{session_id}/runs/{run_id}/cursor", + "", + "other-workspace", + ), + ) + .await; + assert_eq!(cursor_status, StatusCode::OK, "{cursor}"); + let payload = json!({"idempotency_key":"laundered-history", "expected_parent_cursor":cursor, "reason":"B inherited A history"}).to_string(); + fork_status = Some( + self::response( + &fixture.app, + signed_in_scope( + "POST", + &format!("/sessions/{}/forks", fixture.parent), + "/sessions/{session_id}/forks", + &payload, + "other-workspace", + ), + ) + .await + .0, + ); + } + let leaked = fixture + .provider + .requests + .lock() + .await + .iter() + .skip(1) + .any(|request| { + request.body["messages"] + .to_string() + .contains("WORKSPACE-A-PRIVATE-HISTORY") + }); + fixture.close().await; + assert!( + status == StatusCode::FORBIDDEN + || parsed + .iter() + .any(|event| event["error_code"] == "session_fork_forbidden"), + "B continued A parent: {status}; B then prepared fork: {fork_status:?}; model received A history: {leaked}" + ); + assert!(!leaked); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "real disposable MatrixOne; invalid ordinary provider scope proof"] +async fn provider_fork_review_rejects_ordinary_session_with_invalid_scope_proof() { + let fixture = ReviewFixture::new().await; + let owner = fixture.parent_cursor["owner_id"].as_str().unwrap(); + let pool = fixture.state.shared_pool.as_ref().unwrap(); + let raw: String = sqlx::query_scalar( + "SELECT CAST(`metadata` AS CHAR) FROM agent_sessions WHERE user_id = ? AND session_id = ?", + ) + .bind(owner) + .bind(&fixture.parent) + .fetch_one(pool.get()) + .await + .unwrap(); + let mut metadata: Value = serde_json::from_str(&raw).unwrap(); + metadata["__astra_provider_scope_proof_v1"] = Value::Null; + sqlx::query("UPDATE agent_sessions SET `metadata` = ? WHERE user_id = ? AND session_id = ?") + .bind(metadata.to_string()) + .bind(owner) + .bind(&fixture.parent) + .execute(pool.get()) + .await + .unwrap(); + let response = fixture + .app + .clone() + .oneshot(signed_http_request( + "POST", + "/chat/stream", + "/chat/stream", + fixture.moi.payload("MISSING-PROOF").to_string(), + )) + .await + .unwrap(); + let status = response.status(); + let body = to_bytes(response.into_body(), 1024 * 1024).await.unwrap(); + let body = String::from_utf8(body.to_vec()).unwrap(); + let requests = fixture.provider.requests.lock().await.len(); + let interaction_path = "/provider-interactions/respond"; + let interaction = json!({"run_id": &fixture.parent_run, "session_id": &fixture.parent, "request_id": "not-created", "cancelled": true}); + let (interaction_status, interaction_body) = self::response( + &fixture.app, + signed_http_request( + "POST", + interaction_path, + interaction_path, + interaction.to_string(), + ), + ) + .await; + fixture.close().await; + assert!( + status == StatusCode::FORBIDDEN + || events(&body) + .iter() + .any(|event| event["error_code"] == "session_fork_forbidden"), + "{status} {body}" + ); + assert_eq!( + requests, 1, + "invalid proof must not start another model request" + ); + assert_eq!( + interaction_status, + StatusCode::FORBIDDEN, + "{interaction_body}" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "real disposable MatrixOne; existing unbound chat is not a fork source"] +async fn provider_fork_review_preserves_unbound_chat_without_claiming_its_history() { + let mut fixture = ReviewFixture::new().await; + let owner = fixture.parent_cursor["owner_id"] + .as_str() + .unwrap() + .to_owned(); + let pool = fixture.state.shared_pool.as_ref().unwrap().clone(); + let raw: String = sqlx::query_scalar( + "SELECT CAST(`metadata` AS CHAR) FROM agent_sessions WHERE user_id = ? AND session_id = ?", + ) + .bind(&owner) + .bind(&fixture.parent) + .fetch_one(pool.get()) + .await + .unwrap(); + // Model the persisted data from before scope binding was recorded. + let mut metadata: Value = serde_json::from_str(&raw).unwrap(); + metadata + .as_object_mut() + .unwrap() + .remove("__astra_provider_scope_proof_v1"); + sqlx::query("UPDATE agent_sessions SET `metadata` = ? WHERE user_id = ? AND session_id = ?") + .bind(metadata.to_string()) + .bind(&owner) + .bind(&fixture.parent) + .execute(pool.get()) + .await + .unwrap(); + let run = fixture + .moi + .chat( + &fixture.app, + "EXISTING-CONVERSATION-CONTINUES", + "REVIEW-ANSWER", + ) + .await; + let (cursor_status, cursor_error) = control( + &fixture.app, + "GET", + &format!("/sessions/{}/runs/{run}/cursor", fixture.parent), + "/sessions/{session_id}/runs/{run_id}/cursor", + "", + ) + .await; + let payload = json!({"idempotency_key":"unbound-history", "expected_parent_cursor":fixture.parent_cursor, "reason":"cannot claim old history"}); + let (fork_status, fork_error) = control( + &fixture.app, + "POST", + &format!("/sessions/{}/forks", fixture.parent), + "/sessions/{session_id}/forks", + &payload.to_string(), + ) + .await; + let stored: String = sqlx::query_scalar( + "SELECT CAST(`metadata` AS CHAR) FROM agent_sessions WHERE user_id = ? AND session_id = ?", + ) + .bind(&owner) + .bind(&fixture.parent) + .fetch_one(pool.get()) + .await + .unwrap(); + let forks: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM session_forks WHERE owner_user_id = ?") + .bind(&owner) + .fetch_one(pool.get()) + .await + .unwrap(); + fixture.close().await; + assert_eq!(cursor_status, StatusCode::FORBIDDEN, "{cursor_error}"); + assert_eq!(fork_status, StatusCode::FORBIDDEN, "{fork_error}"); + assert!( + serde_json::from_str::(&stored) + .unwrap() + .get("__astra_provider_scope_proof_v1") + .is_none() + ); + assert_eq!( + forks, 0, + "a later Run must not make unbound history forkable" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "real disposable MatrixOne; first provider chat creates a scoped session"] +async fn provider_fork_review_scopes_new_ordinary_chat_session() { + let fixture = ReviewFixture::new().await; + let mut first = fixture.moi.payload("NEW-ORDINARY-HISTORY"); + first.as_object_mut().unwrap().remove("session_id"); + let (status, body) = post(&fixture.app, "/chat/stream", &first).await; + assert_eq!(status, StatusCode::OK, "{body}"); + let parsed = events(&body); + assert!( + parsed.iter().any(|event| event["type"] == "run_finished"), + "{body}" + ); + let session_id = parsed + .iter() + .find(|event| event["type"] == "session_info") + .and_then(|event| event["session_id"].as_str()) + .expect("new session id"); + let mut other = fixture.moi.payload("OTHER-WORKSPACE"); + other["session_id"] = json!(session_id); + let response = fixture + .app + .clone() + .oneshot(signed_in_scope( + "POST", + "/chat/stream", + "/chat/stream", + &other.to_string(), + "other-workspace", + )) + .await + .unwrap(); + let status = response.status(); + let body = to_bytes(response.into_body(), 1024 * 1024).await.unwrap(); + let body = String::from_utf8(body.to_vec()).unwrap(); + let requests = fixture.provider.requests.lock().await.len(); + fixture.close().await; + assert!( + status == StatusCode::FORBIDDEN + || events(&body) + .iter() + .any(|event| event["error_code"] == "session_fork_forbidden"), + "{status} {body}" + ); + assert_eq!(requests, 2, "other workspace must not invoke the model"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "real disposable MatrixOne; PR review prepare cancellation regression"] +async fn provider_fork_review_cancels_interrupted_prepare() { + let fixture = ReviewFixture::new().await; + let payload = json!({"idempotency_key":"cancel-prepare", "expected_parent_cursor":fixture.parent_cursor, "reason":"cancel pending continuation"}).to_string(); + let (status, manifest) = control( + &fixture.app, + "POST", + &format!("/sessions/{}/forks", fixture.parent), + "/sessions/{session_id}/forks", + &payload, + ) + .await; + assert_eq!(status, StatusCode::OK, "{manifest}"); + let fork_id = manifest["fork_id"].as_str().unwrap(); + let (status, value) = control( + &fixture.app, + "POST", + &format!("/sessions/{}/forks/{fork_id}/abort", fixture.parent), + "/sessions/{session_id}/forks/{fork_id}/abort", + r#"{"detail":"request cancelled","grace_seconds":1}"#, + ) + .await; + assert_eq!( + status, + StatusCode::OK, + "provider can cancel a prepared child: {value}" + ); + assert_eq!(value["state"], "aborted"); + let path = format!("/sessions/{}/forks", fixture.parent); + let abort_path = format!("{path}/{fork_id}/abort"); + assert_eq!( + control( + &fixture.app, + "POST", + &abort_path, + "/sessions/{session_id}/forks/{fork_id}/abort", + r#"{"detail":"request cancelled","grace_seconds":1}"# + ) + .await + .0, + StatusCode::OK + ); + let (status, retry) = control( + &fixture.app, + "POST", + &path, + "/sessions/{session_id}/forks", + &payload, + ) + .await; + assert_eq!(status, StatusCode::GONE, "{retry}"); + assert_eq!(retry["error_code"], "session_fork_aborted"); + let activate_path = format!("{path}/{fork_id}/activate-for-continuation"); + assert_eq!( + control( + &fixture.app, + "POST", + &activate_path, + "/sessions/{session_id}/forks/{fork_id}/activate-for-continuation", + "{}" + ) + .await + .0, + StatusCode::GONE + ); + + let pool = fixture.state.shared_pool.as_ref().unwrap(); + let owner = fixture.parent_cursor["owner_id"].as_str().unwrap(); + let child_id = manifest["child_key"]["session_id"].as_str().unwrap(); + let coordinator = DatabaseSessionForkCoordinator::new( + pool.clone(), + Arc::new(DatabaseSessionContextCoordinator::new(pool.clone())), + ); + // Model a completely lost prepare response. Its persisted pin deadline is + // the recovery authority; a caller need not retain the returned fork ID. + let lost_payload = json!({"idempotency_key":"lost-prepare", "expected_parent_cursor":fixture.parent_cursor, "reason":"lost response"}).to_string(); + let (status, lost) = control( + &fixture.app, + "POST", + &path, + "/sessions/{session_id}/forks", + &lost_payload, + ) + .await; + assert_eq!(status, StatusCode::OK, "{lost}"); + sqlx::query("UPDATE conversation_manifest_pins SET grace_expires_at_ms = 0 WHERE owner_user_id = ? AND pin_id IN (?, ?)") + .bind(owner).bind(fork_id).bind(lost["fork_id"].as_str().unwrap()).execute(pool.get()).await.unwrap(); + assert_eq!( + coordinator + .abort_expired_provider_prepares(100) + .await + .unwrap(), + 1 + ); + let (status, expired) = control( + &fixture.app, + "POST", + &path, + "/sessions/{session_id}/forks", + &lost_payload, + ) + .await; + assert_eq!(status, StatusCode::GONE, "{expired}"); + assert_eq!(expired["error_code"], "session_fork_aborted"); + + sqlx::query("UPDATE agent_session_lifecycle_fences SET delete_requested_at = DATE_SUB(NOW(6), INTERVAL 1 SECOND) WHERE user_id = ? AND delete_requested_at IS NOT NULL") + .bind(owner).execute(pool.get()).await.unwrap(); + sqlx::query("UPDATE conversation_manifest_pins SET grace_expires_at_ms = 0 WHERE owner_user_id = ? AND pin_state = 'grace'") + .bind(owner).execute(pool.get()).await.unwrap(); + let swept = astra_services::runtime_maintenance::maintain_runtime_storage( + pool, + Some(&coordinator), + &astra_services::runtime_maintenance::RuntimeMaintenancePolicy { + batch_limit: 100, + delete_intent_grace_secs: 0, + }, + &mut Default::default(), + ) + .await; + assert!( + swept.cleanup_errors.is_empty(), + "{:?}", + swept.cleanup_errors + ); + let children: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM agent_sessions WHERE user_id = ? AND session_id IN (?, ?)", + ) + .bind(owner) + .bind(child_id) + .bind(lost["child_key"]["session_id"].as_str().unwrap()) + .fetch_one(pool.get()) + .await + .unwrap(); + assert_eq!(children, 0, "both cancelled children must be reclaimed"); + let pins: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM conversation_manifest_pins WHERE owner_user_id = ?", + ) + .bind(owner) + .fetch_one(pool.get()) + .await + .unwrap(); + assert_eq!( + pins, 0, + "cancelled preparations cannot pin history indefinitely" + ); + assert_eq!( + control( + &fixture.app, + "POST", + &path, + "/sessions/{session_id}/forks", + &payload + ) + .await + .0, + StatusCode::GONE + ); + assert_eq!( + control( + &fixture.app, + "POST", + &abort_path, + "/sessions/{session_id}/forks/{fork_id}/abort", + r#"{"detail":"request cancelled","grace_seconds":1}"# + ) + .await + .0, + StatusCode::OK + ); + // Active children are never cancelled by this recovery owner. + let (active, active_child) = fork( + &fixture.app, + &fixture.parent, + &fixture.parent_cursor, + "active-control", + ) + .await; + let active_id = active["manifest"]["fork_id"].as_str().unwrap(); + assert_eq!( + control( + &fixture.app, + "POST", + &format!("{path}/{active_id}/abort"), + "/sessions/{session_id}/forks/{fork_id}/abort", + r#"{"detail":"must not delete active","grace_seconds":1}"# + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!( + fixture + .sessions + .get_session(active_child.clone(), owner.into()) + .await + .unwrap() + .status, + "active" + ); + fixture + .sessions + .delete_session(active_child, owner.into()) + .await + .unwrap(); + fixture + .sessions + .delete_session(fixture.parent.clone(), owner.into()) + .await + .unwrap(); + coordinator + .collect_orphaned_fork_records(100) + .await + .unwrap(); + let receipts: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM session_forks WHERE owner_user_id = ?") + .bind(owner) + .fetch_one(pool.get()) + .await + .unwrap(); + assert_eq!( + receipts, 0, + "parent deletion must release its cancelled-fork receipts" + ); + fixture.close().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "real disposable MatrixOne; PR review fork quota regression"] +async fn provider_fork_review_charges_creation_once() { + let fixture = ReviewFixture::new().await; + let owner = fixture.parent_cursor["owner_id"].as_str().unwrap(); + sqlx::query("INSERT INTO resource_limits (user_id, max_sessions_per_day) VALUES (?, 2)") + .bind(owner) + .execute(fixture.state.shared_pool.as_ref().unwrap().get()) + .await + .unwrap(); + fork( + &fixture.app, + &fixture.parent, + &fixture.parent_cursor, + "quota-child", + ) + .await; + let payload = json!({"idempotency_key":"over-quota", "expected_parent_cursor":fixture.parent_cursor, "reason":"second fork must exceed daily limit"}).to_string(); + let (status, value) = control( + &fixture.app, + "POST", + &format!("/sessions/{}/forks", fixture.parent), + "/sessions/{session_id}/forks", + &payload, + ) + .await; + let charged = fixture + .state + .resource_governor + .get_usage(owner) + .await + .sessions_created; + fixture.close().await; + assert_eq!( + status, + StatusCode::TOO_MANY_REQUESTS, + "fork must share the configured daily limit: {value}" + ); + assert_eq!(charged, 2, "prepare replay cannot charge creation twice"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "real disposable MatrixOne; provider-signed Session fork contract"] +async fn provider_signed_forks_continue_from_exact_history_without_inherited_authority() { + assert_eq!(std::env::var("ASTRA_TEST_DB_IT").as_deref(), Ok("1")); + let mut settings = AppSettings::from_explicit_env().expect("explicit test settings"); + let database = format!("astra_test_moi_fork_{}", Uuid::new_v4().simple()); + eprintln!("MOI fork fixture database: {database}"); + settings.matrixone.database = "mysql".into(); + let mut admin = MySqlConnection::connect(&settings.matrixone.database_url_with_password()) + .await + .unwrap(); + sqlx::query(sqlx::AssertSqlSafe(format!("CREATE DATABASE `{database}`"))) + .execute(&mut admin) + .await + .unwrap(); + settings.matrixone.database = database.clone(); + settings.provider_request_auth = vec![ProviderRequestAuthConfig { + provider: "moi".into(), + auth_type: "hmac".into(), + key: KEY.into(), + }]; + settings.external_providers.clear(); + let gateway_state = Arc::new(Mutex::new(GatewayState::default())); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + settings.memoria.base_url = endpoint.clone(); + settings.edge_token_auth = astra_core::EdgeTokenAuthConfig { + key: super::fork_auth_compat::USER_TOKEN_KEY.into(), + check_endpoint: format!("{endpoint}/check"), + }; + let gateway_app = gateway_router(gateway_state.clone()); + let gateway_task = tokio::spawn(async move { + axum::serve(listener, gateway_app).await.unwrap(); + }); + let session_settings = settings.matrixone.clone(); + let state = build_server_state(settings) + .await + .expect("fresh Astra bootstrap"); + let app = build_app(state.clone()); + let mut moi = MoiClient::new(false, endpoint, "fork-contract", gateway_state.clone()); + moi.repair_session(&app).await; + moi.repair_bindings(&app).await; + moi.history.clear(); + let parent = moi.session.clone().unwrap(); + let provider = ProviderGateway::start(vec![ProviderScript::new( + "fork turns", + |r| r.path == "/v1/chat/completions", + [ + "FIRST-ANSWER", + "SECOND-ANSWER", + "THIRD-ANSWER", + "CHILD-ANSWER", + "PARENT-ANSWER", + "SIBLING-ANSWER", + "CHILD-NEXT", + "NESTED-ANSWER", + "AFTER-CHILD-DELETE-ANSWER", + "AFTER-ANCESTOR-DELETE-ANSWER", + ] + .into_iter() + .map(answer) + .collect(), + )]) + .await; + moi.model_endpoint = Some(format!("{}/v1/chat/completions", provider.base_url)); + let first = moi.chat(&app, "FIRST-FACT", "FIRST-ANSWER").await; + let first_cursor = cursor(&app, &parent, &first).await; + assert_eq!(first_cursor["completed_turn"], 1); + let second = moi.chat(&app, "SECOND-FACT", "SECOND-ANSWER").await; + moi.chat(&app, "THIRD-FACT", "THIRD-ANSWER").await; + assert_eq!( + cursor(&app, &parent, &first).await, + first_cursor, + "historical run cannot resolve to latest head" + ); + let path = format!("/sessions/{parent}/forks"); + let payload = json!({"idempotency_key":"tamper-control","expected_parent_cursor":first_cursor,"reason":"test"}).to_string(); + let mut tampered_body = signed_http_request( + "POST", + &path, + "/sessions/{session_id}/forks", + payload.clone(), + ); + *tampered_body.body_mut() = Body::from(payload.replace("tamper-control", "tampered")); + assert_eq!( + response(&app, tampered_body).await.0, + StatusCode::UNAUTHORIZED + ); + let cursor_path = format!("/sessions/{parent}/runs/{first}/cursor"); + let cursor_route = "/sessions/{session_id}/runs/{run_id}/cursor"; + let mut tampered_path = signed_http_request("GET", &cursor_path, cursor_route, String::new()); + *tampered_path.uri_mut() = format!("/sessions/{parent}/runs/{second}/cursor") + .parse() + .unwrap(); + assert_eq!( + response(&app, tampered_path).await.0, + StatusCode::UNAUTHORIZED + ); + let request = signed_http_request("GET", &cursor_path, cursor_route, String::new()); + let headers = request.headers().clone(); + assert_eq!(response(&app, request).await.0, StatusCode::OK); + let mut replay = Request::builder() + .method("GET") + .uri(&cursor_path) + .body(Body::empty()) + .unwrap(); + *replay.headers_mut() = headers; + assert_eq!( + response(&app, replay).await.0, + StatusCode::UNAUTHORIZED, + "request ID replay is not an idempotency key" + ); + let (active, child) = fork(&app, &parent, &first_cursor, "fork-child").await; + let mut mismatch: Value = serde_json::from_str(&payload).unwrap(); + mismatch["idempotency_key"] = json!("fork-child"); + mismatch["reason"] = json!("different intent"); + let (status, error) = control( + &app, + "POST", + &path, + "/sessions/{session_id}/forks", + &mismatch.to_string(), + ) + .await; + assert_eq!(status, StatusCode::CONFLICT); + assert_eq!(error["error_code"], "session_fork_idempotency_mismatch"); + let (status, error) = control( + &app, + "GET", + &format!("/sessions/{child}/runs/{first}/cursor"), + "/sessions/{session_id}/runs/{run_id}/cursor", + "", + ) + .await; + assert_eq!(status, StatusCode::NOT_FOUND); + assert_eq!(error["error_code"], "session_fork_not_found"); + let fork_id = active["manifest"]["fork_id"].as_str().unwrap(); + // A genuinely authenticated legacy credential for the same MOI subject + // still belongs to a different owner namespace; it cannot reach this fork + // through legacy Session reads or controls. + let now = chrono::Utc::now().timestamp(); + let legacy_token = super::fork_auth_compat::token(json!({ + "sub":"moi-contract-user", "workspace_id":"other-workspace", + "iss":"moi-core", "iat":now, "exp":now+3600, "purpose":"runtime" + })); + let own_list = Request::builder() + .uri("/sessions") + .header("authorization", format!("Bearer {legacy_token}")) + .body(Body::empty()) + .unwrap(); + let (legacy_status, legacy_sessions) = response(&app, own_list).await; + assert_eq!( + legacy_status, + StatusCode::OK, + "legacy credential is valid: {legacy_sessions}" + ); + assert!( + legacy_sessions["sessions"].as_array().unwrap().is_empty(), + "legacy cannot list HMAC-owned Sessions" + ); + for (method, target, body) in [ + ("GET", format!("/sessions/{parent}/forks"), String::new()), + ( + "POST", + format!("/sessions/{child}/context/segments:batchGet"), + json!({"manifest_root": first_cursor["canonical_root_hash"]}).to_string(), + ), + ("POST", format!("/sessions/{child}/cancel"), String::new()), + ("DELETE", format!("/sessions/{child}"), String::new()), + ( + "POST", + format!("/sessions/{parent}/forks/{fork_id}/activate"), + json!({"placement":"server","writer_ttl_seconds":30}).to_string(), + ), + ] { + let request = Request::builder() + .method(method) + .uri(target) + .header("content-type", "application/json") + .header("authorization", format!("Bearer {legacy_token}")) + .body(Body::from(body)) + .unwrap(); + let (status, error) = response(&app, request).await; + assert_eq!( + status, + StatusCode::NOT_FOUND, + "legacy cannot own HMAC fork: {error}" + ); + } + for (method, path, route, body) in [ + ( + "GET", + format!("/sessions/{parent}/runs/{first}/cursor"), + "/sessions/{session_id}/runs/{run_id}/cursor", + String::new(), + ), + ( + "POST", + format!("/sessions/{parent}/forks"), + "/sessions/{session_id}/forks", + payload.clone(), + ), + ( + "POST", + format!("/sessions/{parent}/forks/{fork_id}/activate-for-continuation"), + "/sessions/{session_id}/forks/{fork_id}/activate-for-continuation", + "{}".into(), + ), + ] { + let (status, error) = response( + &app, + signed_in_scope(method, &path, route, &body, "other-workspace"), + ) + .await; + assert_eq!( + status, + StatusCode::FORBIDDEN, + "cross-scope request: {error}" + ); + assert_eq!(error["error_code"], "session_fork_forbidden"); + } + let missing_run_path = format!("/sessions/{parent}/runs/{}/cursor", Uuid::new_v4()); + let (status, error) = control( + &app, + "GET", + &missing_run_path, + "/sessions/{session_id}/runs/{run_id}/cursor", + "", + ) + .await; + assert_eq!(status, StatusCode::NOT_FOUND); + assert_eq!(error["error_code"], "session_fork_not_found"); + for (method, path, route, body) in [ + ( + "GET", + format!("/sessions/{parent}/forks"), + "/sessions/{session_id}/forks", + "", + ), + ( + "GET", + format!("/sessions/{parent}/forks/{fork_id}"), + "/sessions/{session_id}/forks/{fork_id}", + "", + ), + ( + "POST", + format!("/sessions/{parent}/forks/{fork_id}/activate"), + "/sessions/{session_id}/forks/{fork_id}/activate", + r#"{"placement":"server","writer_ttl_seconds":30}"#, + ), + ] { + let (status, error) = control(&app, method, &path, route, body).await; + assert_eq!( + status, + StatusCode::UNAUTHORIZED, + "legacy control must retain its existing authentication: {error}" + ); + } + let (_, sibling) = fork(&app, &parent, &first_cursor, "fork-sibling").await; + assert_ne!(child, sibling); + assert_eq!( + provider.requests.lock().await.len(), + 3, + "fork creates no model calls" + ); + moi.session = Some(child.clone()); + let child_run = moi.chat(&app, "CHILD-FACT", "CHILD-ANSWER").await; + let child_cursor = cursor(&app, &child, &child_run).await; + assert_eq!( + child_cursor["completed_turn"], 2, + "child canonical cursor starts after its shared prefix" + ); + moi.session = Some(parent.clone()); + moi.chat(&app, "PARENT-FACT", "PARENT-ANSWER").await; + moi.session = Some(sibling.clone()); + moi.chat(&app, "SIBLING-FACT", "SIBLING-ANSWER").await; + moi.session = Some(child.clone()); + let next = moi.chat(&app, "CHILD-NEXT-FACT", "CHILD-NEXT").await; + assert_eq!(cursor(&app, &child, &next).await["completed_turn"], 3); + let id = active["manifest"]["fork_id"].as_str().unwrap(); + let (status, replay) = control( + &app, + "POST", + &format!("/sessions/{parent}/forks/{id}/activate-for-continuation"), + "/sessions/{session_id}/forks/{fork_id}/activate-for-continuation", + "{}", + ) + .await; + assert_eq!(status, StatusCode::OK); + assert_eq!( + replay["child_head"]["cursor"]["completed_turn"], 3, + "activation replay must not rewind child" + ); + let (_, nested) = fork(&app, &child, &child_cursor, "fork-nested").await; + moi.session = Some(nested.clone()); + let nested_run = moi.chat(&app, "NESTED-FACT", "NESTED-ANSWER").await; + assert_eq!( + cursor(&app, &nested, &nested_run).await["completed_turn"], + 3 + ); + // Each chat consumes terminal + [DONE] and its cursor proves canonical + // settlement. drain_background_runs now shuts admission down, so it is + // reserved for final fixture teardown after all continuation assertions. + + // Persisted adversarial states must never resolve to the Session's latest head. + let pool = state.shared_pool.as_ref().unwrap().get(); + let owner = first_cursor["owner_id"].as_str().unwrap(); + let first_path = format!("/sessions/{parent}/runs/{first}/cursor"); + let cursor_route = "/sessions/{session_id}/runs/{run_id}/cursor"; + sqlx::query( + "UPDATE agent_runs SET parent_run_id = ?, depth = 1 WHERE user_id = ? AND run_id = ?", + ) + .bind("delegating-run") + .bind(owner) + .bind(&first) + .execute(pool) + .await + .unwrap(); + let (status, value) = control(&app, "GET", &first_path, cursor_route, "").await; + assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY); + assert_eq!(value["error_code"], "session_fork_cursor_unsupported"); + sqlx::query("UPDATE agent_runs SET parent_run_id = NULL, depth = 0, status = 'running' WHERE user_id = ? AND run_id = ?") + .bind(owner).bind(&first).execute(pool).await.unwrap(); + let (status, value) = control(&app, "GET", &first_path, cursor_route, "").await; + assert_eq!(status, StatusCode::CONFLICT); + assert_eq!(value["error_code"], "session_fork_cursor_not_ready"); + let (status, value) = response( + &app, + signed_in_scope("GET", &first_path, cursor_route, "", "other-workspace"), + ) + .await; + assert_eq!( + status, + StatusCode::FORBIDDEN, + "scope is checked before readiness" + ); + assert_eq!(value["error_code"], "session_fork_forbidden"); + sqlx::query("UPDATE agent_runs SET status = 'completed' WHERE user_id = ? AND run_id = ?") + .bind(owner) + .bind(&first) + .execute(pool) + .await + .unwrap(); + sqlx::query("UPDATE session_transcript_items SET canonical_completed_turn = NULL, canonical_conversation_seq = NULL, canonical_root_hash = NULL WHERE user_id = ? AND run_id = ?") + .bind(owner).bind(&first).execute(pool).await.unwrap(); + let (status, value) = control(&app, "GET", &first_path, cursor_route, "").await; + assert_eq!(status, StatusCode::CONFLICT); + assert_eq!(value["error_code"], "session_fork_cursor_not_ready"); + sqlx::query("UPDATE session_transcript_items SET canonical_completed_turn = ?, canonical_conversation_seq = ?, canonical_root_hash = ? WHERE user_id = ? AND run_id = ?") + .bind(first_cursor["completed_turn"].as_i64().unwrap()).bind(first_cursor["conversation_seq"].as_i64().unwrap()) + .bind(first_cursor["canonical_root_hash"].as_str().unwrap()).bind(owner).bind(&first).execute(pool).await.unwrap(); + sqlx::query("UPDATE conversation_manifest_nodes SET reachable = 0 WHERE isolation_domain = 'server' AND owner_user_id = ? AND session_id = ? AND manifest_root = ?") + .bind(owner).bind(&parent).bind(first_cursor["canonical_root_hash"].as_str().unwrap()).execute(pool).await.unwrap(); + let (status, value) = control(&app, "GET", &first_path, cursor_route, "").await; + assert_eq!(status, StatusCode::CONFLICT); + assert_eq!(value["error_code"], "session_fork_cursor_unavailable"); + sqlx::query("UPDATE conversation_manifest_nodes SET reachable = 1 WHERE isolation_domain = 'server' AND owner_user_id = ? AND session_id = ? AND manifest_root = ?") + .bind(owner).bind(&parent).bind(first_cursor["canonical_root_hash"].as_str().unwrap()).execute(pool).await.unwrap(); + assert_eq!(cursor(&app, &parent, &first).await, first_cursor); + // A scoped Session can contain an older completed Run with no canonical + // history. Session authorization must not invent its missing fork cursor. + let (status, legacy) = control( + &app, + "POST", + "/sessions", + "/sessions", + &json!({"client_session_ref":format!("legacy-{}", Uuid::new_v4())}).to_string(), + ) + .await; + assert!(status == StatusCode::OK || status == StatusCode::CREATED); + let legacy_session = legacy["session_id"].as_str().unwrap(); + let legacy_run = Uuid::new_v4().to_string(); + sqlx::query("INSERT INTO agent_runs (run_id, user_id, session_id, root_run_id, ancestor_path, status) VALUES (?, ?, ?, ?, ?, 'completed')") + .bind(&legacy_run).bind(owner).bind(legacy_session).bind(&legacy_run).bind(&legacy_run) + .execute(pool).await.unwrap(); + let legacy_path = format!("/sessions/{legacy_session}/runs/{legacy_run}/cursor"); + let (status, value) = control(&app, "GET", &legacy_path, cursor_route, "").await; + assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY); + assert_eq!(value["error_code"], "session_fork_cursor_unsupported"); + // A diagnostic event cannot supply the missing canonical commit either. + let diagnostic_started = json!({"event_type":"run_started", "data":{ + "provider_run_owner":{"provider_id":"moi","provider_scope_id":"moi-contract-workspace"} + }}) + .to_string(); + let diagnostic_hash = format!("{:x}", Sha256::digest(diagnostic_started.as_bytes())); + let diagnostic_insert = sqlx::query("INSERT INTO agent_run_events (id, run_id, event_idx, user_id, session_id, event_type, event_id, event_hash, payload_json) VALUES (?, ?, 0, ?, ?, 'run_started', ?, ?, ?)") + .bind(Uuid::new_v4().to_string()).bind(&legacy_run).bind(owner).bind(legacy_session) + .bind(Uuid::new_v4().to_string()).bind(diagnostic_hash).bind(diagnostic_started) + .execute(pool).await.unwrap(); + assert_eq!(diagnostic_insert.rows_affected(), 1); + let (status, value) = control(&app, "GET", &legacy_path, cursor_route, "").await; + assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY); + assert_eq!(value["error_code"], "session_fork_cursor_unsupported"); + + // The middle Session is ordinary mutable state; its descendants retain + // immutable history without retaining its live head or execution authority. + let shared_pool = state.shared_pool.as_ref().unwrap().clone(); + let context = Arc::new(DatabaseSessionContextCoordinator::new(shared_pool.clone())); + let fork_owner = DatabaseSessionForkCoordinator::new(shared_pool.clone(), context.clone()); + let sessions = DatabaseSessionService::new(session_settings).with_pool(shared_pool); + let nested_key: SessionKeyV1 = serde_json::from_value(json!({ + "schema_version": 1, "isolation_domain": "server", "owner_user_id": owner, + "session_id": nested, "branch_id": "main" + })) + .unwrap(); + let nested_head = context.load_head(&nested_key).await.unwrap().unwrap(); + let exact_nested = context.materialize(&nested_head).await.unwrap(); + sessions + .delete_session(child.clone(), owner.into()) + .await + .unwrap(); + let middle_heads: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM session_context_heads WHERE owner_user_id = ? AND session_id = ?", + ) + .bind(owner) + .bind(&child) + .fetch_one(pool) + .await + .unwrap(); + assert_eq!( + middle_heads, 0, + "deletion must remove the middle live context head" + ); + assert_eq!( + context + .materialize(&nested_head) + .await + .expect("nested history survives middle Session deletion"), + exact_nested, + "deleting a middle Session must preserve the exact descendant prefix" + ); + let middle_key: SessionKeyV1 = + serde_json::from_value(active["manifest"]["child_key"].clone()).unwrap(); + let retained_manifest: astra_turn_types::SessionForkManifestV1 = + serde_json::from_value(active["manifest"].clone()).unwrap(); + assert_eq!( + context.load_fork_prefix(&middle_key).await.unwrap(), + Some(retained_manifest.shared_prefix()) + ); + for wrong_key in [ + SessionKeyV1 { + owner_user_id: "another-owner".into(), + ..middle_key.clone() + }, + SessionKeyV1 { + isolation_domain: "another-domain".into(), + ..middle_key.clone() + }, + SessionKeyV1 { + branch_id: "another-branch".into(), + ..middle_key.clone() + }, + ] { + assert!( + context + .load_fork_prefix(&wrong_key) + .await + .unwrap() + .is_none(), + "retained prefix must not escape the complete child key" + ); + } + assert_eq!( + sessions + .get_session(child.clone(), owner.into()) + .await + .unwrap_err() + .0, + StatusCode::NOT_FOUND + ); + let middle_runs: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM agent_runs WHERE user_id = ? AND session_id = ?") + .bind(owner) + .bind(&child) + .fetch_one(pool) + .await + .unwrap(); + assert_eq!( + middle_runs, 0, + "retaining immutable lineage must not retain execution runs" + ); + let incoming_pin: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM conversation_manifest_pins WHERE owner_user_id = ? AND pin_id = ? AND pin_state = 'active'") + .bind(owner).bind(&retained_manifest.fork_id).fetch_one(pool).await.unwrap(); + assert_eq!( + incoming_pin, 1, + "the upstream retention pin must survive middle deletion" + ); + moi.session = Some(nested.clone()); + let after_middle = moi + .chat(&app, "AFTER-CHILD-DELETE-FACT", "AFTER-CHILD-DELETE-ANSWER") + .await; + cursor(&app, &nested, &after_middle).await; + + // Release the independent sibling before deleting A. Otherwise its A pin + // could hide a broken ancestor pin on the B -> C lineage. + sessions + .delete_session(sibling.clone(), owner.into()) + .await + .unwrap(); + let after_middle_head = context.load_head(&nested_key).await.unwrap().unwrap(); + let exact_after_middle = context.materialize(&after_middle_head).await.unwrap(); + sessions + .delete_session(parent.clone(), owner.into()) + .await + .unwrap(); + for _ in 0..4 { + fork_owner.collect_orphaned_fork_records(100).await.unwrap(); + fork_owner + .collect_unpinned_orphan_manifests(100) + .await + .unwrap(); + } + assert_eq!( + sessions + .get_session(parent.clone(), owner.into()) + .await + .unwrap_err() + .0, + StatusCode::NOT_FOUND + ); + let after_middle_history = context + .materialize(&after_middle_head) + .await + .expect("nested history survives ancestor deletion and GC"); + assert_eq!( + after_middle_history, exact_after_middle, + "ancestor deletion and GC preserve the complete exact child history" + ); + let after_ancestor = moi + .chat( + &app, + "AFTER-ANCESTOR-DELETE-FACT", + "AFTER-ANCESTOR-DELETE-ANSWER", + ) + .await; + cursor(&app, &nested, &after_ancestor).await; + + let requests = provider.requests.lock().await; + for (index, included, excluded) in [ + ( + 3, + vec!["FIRST-FACT", "FIRST-ANSWER", "CHILD-FACT"], + vec!["SECOND-FACT", "THIRD-FACT"], + ), + ( + 4, + vec!["FIRST-FACT", "SECOND-FACT", "THIRD-FACT", "PARENT-FACT"], + vec!["CHILD-FACT"], + ), + ( + 5, + vec!["FIRST-FACT", "SIBLING-FACT"], + vec!["SECOND-FACT", "CHILD-FACT", "PARENT-FACT"], + ), + ( + 6, + vec!["FIRST-FACT", "CHILD-FACT", "CHILD-NEXT-FACT"], + vec!["SIBLING-FACT", "PARENT-FACT"], + ), + ( + 7, + vec!["FIRST-FACT", "CHILD-FACT", "NESTED-FACT"], + vec!["CHILD-NEXT-FACT", "SIBLING-FACT", "PARENT-FACT"], + ), + ( + 8, + vec![ + "FIRST-FACT", + "CHILD-FACT", + "NESTED-FACT", + "AFTER-CHILD-DELETE-FACT", + ], + vec!["CHILD-NEXT-FACT", "SIBLING-FACT", "PARENT-FACT"], + ), + ( + 9, + vec![ + "FIRST-FACT", + "CHILD-FACT", + "NESTED-FACT", + "AFTER-CHILD-DELETE-FACT", + "AFTER-ANCESTOR-DELETE-FACT", + ], + vec!["CHILD-NEXT-FACT", "SIBLING-FACT", "PARENT-FACT"], + ), + ] { + let messages = requests[index].body["messages"].to_string(); + for value in included { + assert!(messages.contains(value), "input {index} missing {value}"); + } + for value in excluded { + assert!(!messages.contains(value), "input {index} leaked {value}"); + } + } + drop(requests); + provider.assert_complete(); + sessions.delete_session(nested, owner.into()).await.unwrap(); + for _ in 0..4 { + fork_owner.collect_orphaned_fork_records(100).await.unwrap(); + fork_owner + .collect_unpinned_orphan_manifests(100) + .await + .unwrap(); + } + for table in [ + "session_forks", + "session_fork_events", + "conversation_manifest_pins", + "conversation_manifest_nodes", + ] { + let remaining: i64 = sqlx::query_scalar(sqlx::AssertSqlSafe(format!( + "SELECT COUNT(*) FROM `{table}` WHERE owner_user_id = ?" + ))) + .bind(owner) + .fetch_one(pool) + .await + .unwrap(); + assert_eq!( + remaining, 0, + "retained lineage must converge after the last descendant is deleted: {table}" + ); + } + assert!(state.drain_background_runs(Duration::from_secs(10)).await); + state.close_database_pools().await; + gateway_task.abort(); + sqlx::query(sqlx::AssertSqlSafe(format!("DROP DATABASE `{database}`"))) + .execute(&mut admin) + .await + .unwrap(); +} diff --git a/crates/runtime/tests/web_agent_e2e.rs b/crates/runtime/tests/web_agent_e2e.rs index cd06ede174..1648388b3c 100644 --- a/crates/runtime/tests/web_agent_e2e.rs +++ b/crates/runtime/tests/web_agent_e2e.rs @@ -3,6 +3,8 @@ //! //! These tests exercise the `/chat/stream` → `ServerAgenticLoopHost` path (NOT the bridge), //! using native loopback provider responses through the real Server execution path. +//! Binding ownership cases enter the sessionless runtime boundary directly; +//! real persistent provider Session admission is covered by `moi_fresh_astra`. //! //! ```text //! cargo test -p astra-runtime --test web_agent_e2e --features e2e-hooks @@ -687,7 +689,7 @@ fn build_test_app_with_agent_bindings( binding_service: Arc, models: Arc, inference: Option<&ProviderObservationsFixture>, -) -> (Router, Arc) { +) -> (AppState, Arc) { init_env(); let base = AppState::new(ServiceInfo::default(), Arc::new(StubHealth)) .with_auth_service(Arc::new(StubAuth)) @@ -710,7 +712,7 @@ fn build_test_app_with_agent_bindings( lifecycle }; let state = base.with_run_lifecycle_service(Arc::new(lifecycle)); - (build_app(state), observer_worker) + (state, observer_worker) } // ── Helpers ────────────────────────────────────────────────────────────────── @@ -1318,6 +1320,37 @@ async fn start_agent_binding_gateway( (format!("http://{address}/capabilities"), calls, server) } +// Binding ownership is exercised through real provider HTTP and canonical +// execution. Persistent Session creation/admission is covered by moi_fresh_astra. +async fn provider_binding_stream_collect(state: &AppState, payload: Value) -> Vec { + let mut headers = HeaderMap::new(); + headers.insert("authorization", PROVIDER_TOKEN.parse().unwrap()); + let principal = StubAuth + .current_principal_for_request( + &headers, + ProviderRequestDescriptor { + method: "POST".into(), + path: "/chat/stream".into(), + route: Some("/chat/stream".into()), + request_id: None, + body_digest: None, + }, + ) + .await + .unwrap(); + let response = astra_runtime::server::provider_test_support::ephemeral_provider_chat_stream( + state, + &principal, + normalize_chat_stream_payload(payload), + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + let body = body::to_bytes(response.into_body(), 16 * 1024 * 1024) + .await + .unwrap(); + parse_sse_events(&String::from_utf8_lossy(&body)) +} + fn agent_binding_chat_payload( foundation_binding_id: &str, extension_binding_id: &str, @@ -1393,7 +1426,7 @@ let inference=ProviderObservationsFixture::default(); let model_endpoint=format!("{}/v1",gateway.base_url); let (app,observer_worker)=build_test_app_with_agent_bindings(binding_service, Arc::new(TestModelService{judgment_base_url:Some(model_endpoint.clone())}),Some(&inference)); - let events = provider_chat_stream_collect( + let events = provider_binding_stream_collect( &app, agent_binding_chat_payload( &foundation_binding_id, @@ -1509,7 +1542,7 @@ async fn multi_agent_binding_http_e2e_reports_exact_missing_binding_id() { None, ); - let events = provider_chat_stream_collect( + let events = provider_binding_stream_collect( &app, agent_binding_chat_payload( &foundation_binding_id, diff --git a/crates/services/src/auth/encryption.rs b/crates/services/src/auth/encryption.rs index 3a819804c2..77d59b213a 100644 --- a/crates/services/src/auth/encryption.rs +++ b/crates/services/src/auth/encryption.rs @@ -1,4 +1,7 @@ -use base64::{Engine as _, engine::general_purpose::URL_SAFE}; +use base64::{ + Engine as _, + engine::general_purpose::{URL_SAFE, URL_SAFE_NO_PAD}, +}; use fernet::Fernet; use hmac::{Hmac, Mac}; use sha2::{Digest, Sha256}; @@ -64,6 +67,58 @@ impl FernetTokenEncryptor { mac.update(value.as_bytes()); format!("{:x}", mac.finalize().into_bytes()) } + + fn provider_session_scope_mac( + &self, + user_id: &str, + session_id: &str, + provider_id: &str, + scope_id: &str, + ) -> HmacSha256 { + let mut mac = + HmacSha256::new_from_slice(&self.digest_key).expect("HMAC accepts a 32-byte key"); + mac.update(b"astra.provider-session-scope.v1\0"); + for value in [user_id, session_id, provider_id, scope_id] { + mac.update(&(value.len() as u64).to_be_bytes()); + mac.update(value.as_bytes()); + } + mac + } + + /// Sign the immutable provider scope of one session with the server key. + pub fn sign_provider_session_scope( + &self, + user_id: &str, + session_id: &str, + provider_id: &str, + scope_id: &str, + ) -> String { + URL_SAFE_NO_PAD.encode( + self.provider_session_scope_mac(user_id, session_id, provider_id, scope_id) + .finalize() + .into_bytes(), + ) + } + + /// Verify a scope proof without trusting its storage or its caller. + pub fn verify_provider_session_scope( + &self, + user_id: &str, + session_id: &str, + provider_id: &str, + scope_id: &str, + proof: &str, + ) -> bool { + if proof.len() != 43 { + return false; + } + let Ok(signature) = URL_SAFE_NO_PAD.decode(proof) else { + return false; + }; + self.provider_session_scope_mac(user_id, session_id, provider_id, scope_id) + .verify_slice(&signature) + .is_ok() + } } fn derive_fernet_key(secret_key: &str) -> String { @@ -139,6 +194,81 @@ mod tests { ); } + #[test] + fn provider_session_scope_proof_binds_every_identity_component() { + let key = make_encryptor("scope-key-one"); + let other_key = make_encryptor("scope-key-two"); + let proof = key.sign_provider_session_scope("user", "session", "provider", "workspace"); + assert!(key.verify_provider_session_scope( + "user", + "session", + "provider", + "workspace", + &proof + )); + assert!(!key.verify_provider_session_scope("user", "session", "provider", "other", &proof)); + assert!(!key.verify_provider_session_scope( + "user", + "other", + "provider", + "workspace", + &proof + )); + assert!(!key.verify_provider_session_scope( + "other", + "session", + "provider", + "workspace", + &proof + )); + assert!(!key.verify_provider_session_scope( + "user", + "session", + "other", + "workspace", + &proof + )); + assert!(!other_key.verify_provider_session_scope( + "user", + "session", + "provider", + "workspace", + &proof + )); + assert!(!key.verify_provider_session_scope( + "user", + "session", + "provider", + "workspace", + "invalid" + )); + assert!(!key.verify_provider_session_scope( + "user", + "session", + "provider", + "workspace", + &"!".repeat(43) + )); + assert!(!key.verify_provider_session_scope( + "user", + "session", + "provider", + "workspace", + &proof[..42] + )); + let mut altered = proof.clone().into_bytes(); + altered[0] = if altered[0] == b'A' { b'B' } else { b'A' }; + assert!(!key.verify_provider_session_scope( + "user", + "session", + "provider", + "workspace", + std::str::from_utf8(&altered).unwrap() + )); + let ambiguous = key.sign_provider_session_scope("us", "ersession", "provider", "workspace"); + assert_ne!(proof, ambiguous); + } + #[test] fn decrypt_invalid_ciphertext_fails() { let enc = make_encryptor("any-key"); diff --git a/crates/services/src/auth/mod.rs b/crates/services/src/auth/mod.rs index 55eb09a6fd..b14cc2c36d 100644 --- a/crates/services/src/auth/mod.rs +++ b/crates/services/src/auth/mod.rs @@ -437,6 +437,25 @@ pub enum EdgeTokenBinding { MissingBinding, } +/// Authentication mechanism established by the verifier for this request. +/// Provider identity alone cannot distinguish existing user tokens from the +/// request-bound HMAC mechanism. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RequestAuthMechanism { + /// Successfully verified through the existing current-user credential path. + KnownLegacy, + /// Successfully verified against this request's method, path, body and nonce. + RequestBoundHmac, + /// An existing custom verifier supplied no explicit mechanism evidence. + Unspecified, +} + +#[derive(Clone, Debug)] +pub struct RequestAuthResult { + pub principal: AuthPrincipal, + pub mechanism: RequestAuthMechanism, +} + #[async_trait] pub trait AuthService: Send + Sync { async fn reauthentication_options( @@ -549,6 +568,22 @@ pub trait AuthService: Send + Sync { self.current_principal(headers).await } + /// Existing custom request verifiers remain authoritative. Their successful + /// result is deliberately unspecified: consumers must not infer legacy + /// admission from a provider principal or retry another verifier. + async fn current_principal_for_request_with_provenance( + &self, + headers: &HeaderMap, + request: ProviderRequestDescriptor, + ) -> Result)> { + self.current_principal_for_request(headers, request) + .await + .map(|principal| RequestAuthResult { + principal, + mechanism: RequestAuthMechanism::Unspecified, + }) + } + /// Fresh authorization from a store-created canonical admission before /// claiming execution. Serialized provenance alone cannot grant a principal. async fn reauthorize_execution_handoff( @@ -1265,6 +1300,76 @@ impl DatabaseAuthService { Ok(()) } + async fn verify_request_principal( + &self, + headers: &HeaderMap, + request: ProviderRequestDescriptor, + ) -> Result)> { + let Some((provider_id, _action, token)) = Self::external_request_auth_headers(headers)? + else { + let token = bearer_token(headers)?; + let principal = if token.starts_with(MOI_USER_TOKEN_PREFIX) { + self.principal_from_edge_token(token, request).await? + } else { + self.current_principal(headers).await? + }; + return Ok(RequestAuthResult { + principal, + mechanism: RequestAuthMechanism::KnownLegacy, + }); + }; + let provider = self.provider_config(&provider_id)?.clone(); + if provider.auth_type != "hmac" { + return Err(provider_request_auth_error( + "Provider request auth configuration is unsupported", + )); + } + let authorized = verify_provider_request_hmac_token( + &provider.provider, + &provider.key, + &token, + &request, + )?; + let request_id = request.request_id.as_deref().unwrap_or_default(); + self.record_provider_request_authorization(&authorized, request_id) + .await?; + let user_id = format!( + "provider_authorized:{}:{}", + authorized.provider_id, authorized.external_subject + ); + let username_len = authorized.external_subject.chars().count(); + if username_len > USERNAME_MAX_LEN { + return Err(provider_request_auth_error( + "Provider request external subject exceeds the Astra username limit", + )); + } + if user_id.chars().count() > USER_ID_MAX_LEN { + return Err(provider_request_auth_error( + "Provider-authorized user principal exceeds the Astra user_id limit", + )); + } + let principal = AuthPrincipal::from_provider_request( + AuthUserRecord { + user_id, + username: authorized.external_subject.clone(), + email: String::new(), + display_name: None, + }, + AuthProviderAuthorizedRequestContext { + provider_id: authorized.provider_id, + external_subject: authorized.external_subject, + provider_scope_id: authorized.provider_scope_id, + request_authorization_id: authorized.request_authorization_id, + // Provider-request (HMAC server-to-server) calls carry no edge binding. + edge_agent_id: None, + }, + ); + Ok(RequestAuthResult { + principal, + mechanism: RequestAuthMechanism::RequestBoundHmac, + }) + } + fn external_request_auth_headers( headers: &HeaderMap, ) -> Result { @@ -2257,60 +2362,17 @@ impl AuthService for DatabaseAuthService { headers: &HeaderMap, request: ProviderRequestDescriptor, ) -> Result)> { - let Some((provider_id, _action, token)) = Self::external_request_auth_headers(headers)? - else { - let token = bearer_token(headers)?; - if token.starts_with(MOI_USER_TOKEN_PREFIX) { - return self.principal_from_edge_token(token, request).await; - } - return self.current_principal(headers).await; - }; - let provider = self.provider_config(&provider_id)?.clone(); - if provider.auth_type != "hmac" { - return Err(provider_request_auth_error( - "Provider request auth configuration is unsupported", - )); - } - let authorized = verify_provider_request_hmac_token( - &provider.provider, - &provider.key, - &token, - &request, - )?; - let request_id = request.request_id.as_deref().unwrap_or_default(); - self.record_provider_request_authorization(&authorized, request_id) - .await?; - let user_id = format!( - "provider_authorized:{}:{}", - authorized.provider_id, authorized.external_subject - ); - let username_len = authorized.external_subject.chars().count(); - if username_len > USERNAME_MAX_LEN { - return Err(provider_request_auth_error( - "Provider request external subject exceeds the Astra username limit", - )); - } - if user_id.chars().count() > USER_ID_MAX_LEN { - return Err(provider_request_auth_error( - "Provider-authorized user principal exceeds the Astra user_id limit", - )); - } - Ok(AuthPrincipal::from_provider_request( - AuthUserRecord { - user_id, - username: authorized.external_subject.clone(), - email: String::new(), - display_name: None, - }, - AuthProviderAuthorizedRequestContext { - provider_id: authorized.provider_id, - external_subject: authorized.external_subject, - provider_scope_id: authorized.provider_scope_id, - request_authorization_id: authorized.request_authorization_id, - // Provider-request (HMAC server-to-server) calls carry no edge binding. - edge_agent_id: None, - }, - )) + self.verify_request_principal(headers, request) + .await + .map(|result| result.principal) + } + + async fn current_principal_for_request_with_provenance( + &self, + headers: &HeaderMap, + request: ProviderRequestDescriptor, + ) -> Result)> { + self.verify_request_principal(headers, request).await } async fn edge_registration_binding( diff --git a/crates/services/src/auth/session.rs b/crates/services/src/auth/session.rs index eefe77e31b..4d271cc07a 100644 --- a/crates/services/src/auth/session.rs +++ b/crates/services/src/auth/session.rs @@ -13,8 +13,10 @@ use chrono::Utc; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use sqlx::{Connection, MySql, QueryBuilder, Row, query}; +use std::sync::Arc; use uuid::Uuid; +pub const PROVIDER_SCOPE_PROOF_KEY: &str = "__astra_provider_scope_proof_v1"; const MAX_SESSION_ACTIVITY_ROWS: u32 = 200; const SESSION_LIST_CURSOR_FILTER_SQL: &str = " AND (activity_at < "; const SESSION_LIST_CURSOR_TIE_SQL: &str = " OR (activity_at = "; @@ -177,6 +179,25 @@ pub trait SessionService: Send + Sync { )) } + /// Create a Session with the configured creation quota. Database-backed + /// implementations reserve daily usage in the Session insert transaction. + async fn create_session_with_resource_quota( + &self, + user_id: String, + request: SessionCreateRequestData, + governor: &dyn crate::resource_governor::ResourceGovernor, + ) -> Result<(SessionRecord, Option), (StatusCode, Json)> + { + crate::resource_governor::enforce_session_create_quota( + &governor.check_session_create(&user_id).await, + )?; + let session = self + .create_session_with_admission_facts(user_id.clone(), request) + .await?; + governor.record_session_created(&user_id).await; + Ok(session) + } + async fn list_sessions( &self, filter: SessionListFilter, @@ -268,6 +289,7 @@ pub struct SessionCreateRequestData { pub struct ProviderSessionCreationIdentity { user_id: String, session_id: String, + scope: crate::runs::ProviderRunOwner, } impl ProviderSessionCreationIdentity { @@ -316,6 +338,10 @@ impl ProviderSessionCreationIdentity { Ok(Self { user_id: principal.user.user_id.clone(), session_id: format!("{:x}", hash.finalize()), + scope: crate::runs::ProviderRunOwner { + provider_id: context.provider_id.clone(), + provider_scope_id: context.provider_scope_id.clone(), + }, }) } } @@ -577,10 +603,24 @@ fn session_activity_entry_from_row( }) } -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct DatabaseSessionService { matrixone: MatrixOneSettings, pool: Option, + scope_encryptor: Option>, +} + +impl std::fmt::Debug for DatabaseSessionService { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DatabaseSessionService") + .field("matrixone", &self.matrixone) + .field("pool", &self.pool) + .field( + "scope_encryptor_configured", + &self.scope_encryptor.is_some(), + ) + .finish() + } } #[derive(Clone, Debug, PartialEq)] @@ -605,12 +645,40 @@ impl DatabaseSessionService { session_id: String, request: SessionCreateRequestData, provider_creation: Option<(String, crate::resource_governor::LimitCheck)>, + reserve_quota: bool, + provider_scope: Option<&crate::runs::ProviderRunOwner>, ) -> Result)> { let pool = self.get_pool().await.map_err(internal_error)?; let title = request .title .unwrap_or_else(|| format!("Session {}", Utc::now().format("%Y-%m-%d %H:%M"))); - let metadata = serde_json::Value::Object(request.metadata.unwrap_or_default()).to_string(); + let mut metadata = request.metadata.unwrap_or_default(); + if metadata.contains_key(PROVIDER_SCOPE_PROOF_KEY) { + return Err(error_response_coded( + StatusCode::BAD_REQUEST, + "reserved session metadata key", + "session_metadata_invalid", + )); + } + if let Some(scope) = provider_scope { + let encryptor = self.scope_encryptor.as_ref().ok_or_else(|| { + error_response_coded( + StatusCode::SERVICE_UNAVAILABLE, + "provider session scope is not configured", + "provider_session_scope_unavailable", + ) + })?; + metadata.insert( + PROVIDER_SCOPE_PROOF_KEY.to_string(), + serde_json::Value::String(encryptor.sign_provider_session_scope( + &user_id, + &session_id, + &scope.provider_id, + &scope.provider_scope_id, + )), + ); + } + let metadata = serde_json::Value::Object(metadata).to_string(); let mut connection = CancellationSafePoolConnection::acquire(&pool) .await @@ -643,6 +711,34 @@ impl DatabaseSessionService { .fetch_session_for_user_with_admission_facts(&mut *tx, &session_id, &user_id) .await? .ok_or_else(|| internal_error("failed to read idempotent session"))?; + // Replaying the same deterministic creation intent may read + // an older unbound Session, but must never backfill its scope. + if let Some(scope) = provider_scope + && record.metadata.contains_key(PROVIDER_SCOPE_PROOF_KEY) + { + let valid = record + .metadata + .get(PROVIDER_SCOPE_PROOF_KEY) + .and_then(serde_json::Value::as_str) + .is_some_and(|proof| { + self.scope_encryptor.as_ref().is_some_and(|encryptor| { + encryptor.verify_provider_session_scope( + &user_id, + &session_id, + &scope.provider_id, + &scope.provider_scope_id, + proof, + ) + }) + }); + if !valid { + return Err(error_response_coded( + StatusCode::FORBIDDEN, + "provider session scope cannot be verified", + "provider_session_scope_unverified", + )); + } + } tx.commit().await.map_err(internal_error)?; connection.release(); return Ok(SessionCreationResult { @@ -654,6 +750,16 @@ impl DatabaseSessionService { crate::resource_governor::enforce_session_create_quota(quota)?; } + if reserve_quota { + let quota = + crate::resource_governor::DatabaseResourceGovernor::reserve_session_creation( + &mut tx, &user_id, + ) + .await + .map_err(internal_error)?; + crate::resource_governor::enforce_session_create_quota("a)?; + } + query( "INSERT INTO agent_sessions \ (session_id, user_id, agent_id, title, status, event_count, created_at, updated_at, last_active_at, `metadata`, provider_creation_hash) \ @@ -695,6 +801,7 @@ impl DatabaseSessionService { Self { matrixone, pool: None, + scope_encryptor: None, } } @@ -824,6 +931,11 @@ impl DatabaseSessionService { self } + pub fn with_scope_encryptor(mut self, encryptor: Arc) -> Self { + self.scope_encryptor = Some(encryptor); + self + } + async fn get_pool(&self) -> Result, sqlx::Error> { crate::require_shared_pool( self.pool.as_ref(), @@ -840,9 +952,16 @@ impl SessionService for DatabaseSessionService { user_id: String, request: SessionCreateRequestData, ) -> Result)> { - self.create_session_record(user_id, Uuid::new_v4().to_string(), request, None) - .await - .map(|result| result.session) + self.create_session_record( + user_id, + Uuid::new_v4().to_string(), + request, + None, + false, + None, + ) + .await + .map(|result| result.session) } async fn create_session_with_admission_facts( @@ -851,9 +970,16 @@ impl SessionService for DatabaseSessionService { request: SessionCreateRequestData, ) -> Result<(SessionRecord, Option), (StatusCode, Json)> { - self.create_session_record(user_id, Uuid::new_v4().to_string(), request, None) - .await - .map(|result| (result.session, Some(result.admission_facts))) + self.create_session_record( + user_id, + Uuid::new_v4().to_string(), + request, + None, + false, + None, + ) + .await + .map(|result| (result.session, Some(result.admission_facts))) } async fn create_provider_session( @@ -869,10 +995,31 @@ impl SessionService for DatabaseSessionService { identity.session_id, request, Some((payload_hash, quota)), + true, + Some(&identity.scope), ) .await } + async fn create_session_with_resource_quota( + &self, + user_id: String, + request: SessionCreateRequestData, + _governor: &dyn crate::resource_governor::ResourceGovernor, + ) -> Result<(SessionRecord, Option), (StatusCode, Json)> + { + self.create_session_record( + user_id, + Uuid::new_v4().to_string(), + request, + None, + true, + None, + ) + .await + .map(|result| (result.session, Some(result.admission_facts))) + } + async fn list_sessions( &self, filter: SessionListFilter, @@ -1049,6 +1196,20 @@ impl SessionService for DatabaseSessionService { )); } + if metadata + .as_ref() + .is_some_and(|value| value.contains_key(PROVIDER_SCOPE_PROOF_KEY)) + || metadata_patch + .as_ref() + .is_some_and(|value| value.contains_key(PROVIDER_SCOPE_PROOF_KEY)) + { + return Err(error_response_coded( + StatusCode::BAD_REQUEST, + "reserved session metadata key", + "session_metadata_invalid", + )); + } + if title.is_none() && metadata.is_none() && metadata_patch.is_none() && status.is_none() { return Ok(existing); } @@ -1097,18 +1258,24 @@ impl SessionService for DatabaseSessionService { ) })?; + let raw = locked + .try_get::("metadata_json") + .map_err(internal_error)?; + let value: serde_json::Value = serde_json::from_str(&raw).map_err(internal_error)?; + let locked_metadata = value + .as_object() + .cloned() + .ok_or_else(|| internal_error("persisted session metadata must be a JSON object"))?; let metadata_patch_for_audit = metadata_patch.clone(); let metadata = match (metadata, metadata_patch) { - (Some(metadata), None) => Some(metadata), + (Some(mut metadata), None) => { + if let Some(proof) = locked_metadata.get(PROVIDER_SCOPE_PROOF_KEY) { + metadata.insert(PROVIDER_SCOPE_PROOF_KEY.to_string(), proof.clone()); + } + Some(metadata) + } (None, Some(patch)) => { - let raw = locked - .try_get::("metadata_json") - .map_err(internal_error)?; - let value: serde_json::Value = - serde_json::from_str(&raw).map_err(internal_error)?; - let mut current = value.as_object().cloned().ok_or_else(|| { - internal_error("persisted session metadata must be a JSON object") - })?; + let mut current = locked_metadata; apply_metadata_patch(&mut current, patch); Some(current) } diff --git a/crates/services/src/lib.rs b/crates/services/src/lib.rs index b4fcc4fe6c..0b3227b90a 100644 --- a/crates/services/src/lib.rs +++ b/crates/services/src/lib.rs @@ -147,9 +147,9 @@ pub use auth::{ ExternalProviderClient, ExternalRequestDescriptor, ExternalRuntimeContextRequestData, ExternalRuntimeContextResponse, FernetTokenEncryptor, HttpExternalProviderClient, ProviderRequestDescriptor, ProviderSessionCreationIdentity, ReauthenticationProofRecord, - ReauthenticationPurpose, ReauthenticationRequestData, SessionCreateRequestData, - SessionCreationResult, SessionListFilter, SessionListRecord, SessionRecord, SessionService, - SessionUpdateRequestData, + ReauthenticationPurpose, ReauthenticationRequestData, RequestAuthMechanism, RequestAuthResult, + SessionCreateRequestData, SessionCreationResult, SessionListFilter, SessionListRecord, + SessionRecord, SessionService, SessionUpdateRequestData, }; pub use context::{ ContextService, DatabaseContextService, SnapshotCreateRequestData, SnapshotListFilter, @@ -293,7 +293,8 @@ pub use session_context_coordinator::{ TransferWriterOutcome, WriterTransferConflictV1, WriterTransferRequestV1, }; pub use session_fork_coordinator::{ - DatabaseSessionForkCoordinator, PrepareSessionForkV1, SessionForkCoordinatorError, + DatabaseSessionForkCoordinator, PrepareSessionForkV1, SessionForkContinuationV1, + SessionForkCoordinatorError, }; pub use session_handoff::{ AttachSessionOutcomeV1, AttachSessionRequestV1, ClaimSessionControllerOutcomeV1, diff --git a/crates/services/src/resource_governor.rs b/crates/services/src/resource_governor.rs index 713bae145d..4b04045b76 100644 --- a/crates/services/src/resource_governor.rs +++ b/crates/services/src/resource_governor.rs @@ -10,7 +10,8 @@ //! //! # Design //! -//! * **Fail-open**: DB errors log a warning and return defaults / proceed. +//! * **Read-only checks fail open**: DB errors log a warning and return defaults. +//! Durable Session creation requires its quota counter to commit atomically. //! * **Atomic counters**: `ON DUPLICATE KEY UPDATE` for race-free increments. //! * **No implicit token price policy**: daily tokens are unlimited unless an //! administrator writes a finite per-user limit; `0` means no cap. @@ -294,6 +295,78 @@ impl DatabaseResourceGovernor { Ok(()) } + /// Reserve one durable Session creation in the caller's transaction. + /// + /// Call this after idempotent replay has been ruled out. The existing daily + /// usage row serializes ordinary and fork creation, and the increment rolls + /// back with the Session if any later operation fails. Resource tables must + /// already have been initialized by the database resource governor. + pub async fn reserve_session_creation( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + user_id: &str, + ) -> Result { + let today = Self::today(); + sqlx::query( + "INSERT INTO resource_usage (user_id, usage_date, sessions_created) + VALUES (?, ?, 0) + ON DUPLICATE KEY UPDATE sessions_created = sessions_created", + ) + .bind(user_id) + .bind(&today) + .execute(&mut **tx) + .await?; + let (sessions_created, tokens_consumed): (i64, i64) = sqlx::query_as( + "SELECT CAST(sessions_created AS SIGNED), CAST(tokens_consumed AS SIGNED) + FROM resource_usage WHERE user_id = ? AND usage_date = ? FOR UPDATE", + ) + .bind(user_id) + .bind(&today) + .fetch_one(&mut **tx) + .await?; + let limits: Option<(i32, i32, i64)> = sqlx::query_as( + "SELECT max_concurrent_sessions, max_sessions_per_day, max_tokens_per_day + FROM resource_limits WHERE user_id = ?", + ) + .bind(user_id) + .fetch_optional(&mut **tx) + .await?; + let (concurrent_limit, daily_limit, token_limit) = limits.unwrap_or(( + ResourceLimits::DEFAULT_MAX_CONCURRENT_SESSIONS as i32, + ResourceLimits::DEFAULT_MAX_SESSIONS_PER_DAY as i32, + ResourceLimits::DEFAULT_MAX_TOKENS_PER_DAY as i64, + )); + let active_sessions = if concurrent_limit > 0 { + sqlx::query_scalar( + "SELECT COUNT(DISTINCT session_id) FROM agent_runs + WHERE user_id = ? AND status IN ('running', 'paused', 'waiting')", + ) + .bind(user_id) + .fetch_one(&mut **tx) + .await? + } else { + 0 + }; + let check = Self::evaluate_session_create_snapshot( + i64::from(concurrent_limit), + i64::from(daily_limit), + token_limit, + sessions_created, + tokens_consumed, + active_sessions, + ); + if check == LimitCheck::Allowed { + sqlx::query( + "UPDATE resource_usage SET sessions_created = sessions_created + 1 + WHERE user_id = ? AND usage_date = ?", + ) + .bind(user_id) + .bind(&today) + .execute(&mut **tx) + .await?; + } + Ok(check) + } + /// Count sessions that currently hold execution capacity. /// /// Web sessions are durable and resumable, so historical/open chat sessions must @@ -618,9 +691,9 @@ impl ResourceGovernor for DatabaseResourceGovernor { 0 }; - // This remains an observational, non-reserving check. Hard concurrent - // enforcement still requires an atomic quota reservation with the - // session insert; a read-only check cannot close that race. + // This remains a non-reserving preflight. Database Session creation + // rechecks through reserve_session_creation and commits the daily + // counter with the Session; a read-only check cannot close that race. Self::evaluate_session_create_snapshot( max_concurrent_sessions, max_sessions_per_day, diff --git a/crates/services/src/runs.rs b/crates/services/src/runs.rs index f4d9104b18..1363ab464d 100644 --- a/crates/services/src/runs.rs +++ b/crates/services/src/runs.rs @@ -60,6 +60,21 @@ pub fn is_run_lifecycle_unconfigured_error(status: StatusCode, error: &ErrorResp #[async_trait] pub trait RunLifecycleService: Send + Sync { + /// Verify the provider's Session scope before exposing or controlling a Run. + /// Implementations without an authoritative scope store fail closed. + async fn authorize_provider_run_scope( + &self, + _user_id: String, + _run_id: String, + _owner: ProviderRunOwner, + ) -> Result<(), (StatusCode, Json)> { + Err(error_response_coded( + StatusCode::SERVICE_UNAVAILABLE, + "Provider session scope cannot be verified", + "provider_session_scope_unavailable", + )) + } + async fn request_permission_mode( &self, _user_id: String, @@ -1147,11 +1162,7 @@ impl RunStartIdempotency { /// This value is runtime-injected and never accepted from an HTTP chat /// payload. Durable provider interactions carry it so callbacks enforce the /// same provider-owned scope. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct ProviderRunOwner { - pub provider_id: String, - pub provider_scope_id: String, -} +pub use astra_turn_types::ProviderRunOwner; /// Facts loaded while authenticating and resolving an existing Session. /// @@ -11042,6 +11053,7 @@ enum CancellationMarkerAttempt { #[derive(Clone)] pub struct DatabaseRunStateStore { pool: SharedPool, + scope_encryptor: Option>, owner_pod_id: String, lease_ttl: Duration, #[cfg(any(test, feature = "e2e-hooks"))] @@ -11675,6 +11687,7 @@ impl DatabaseRunStateStore { pub fn new(pool: SharedPool) -> Self { Self { pool, + scope_encryptor: None, owner_pod_id: default_owner_pod_id(), lease_ttl: Self::DEFAULT_LEASE_TTL, #[cfg(any(test, feature = "e2e-hooks"))] @@ -11744,6 +11757,14 @@ impl DatabaseRunStateStore { } } + pub fn with_scope_encryptor( + mut self, + encryptor: std::sync::Arc, + ) -> Self { + self.scope_encryptor = Some(encryptor); + self + } + async fn hydrate_agent_recovery_events( &self, user_id: &str, @@ -15247,6 +15268,24 @@ impl DatabaseRunStateStore { if !admitted { return Err(RunAdmissionError::SessionBusy); } + let provider_owner = record + .opening + .as_ref() + .and_then(|opening| opening.events().first()) + .and_then(|event| event.pointer("/data/provider_run_owner")) + .map(|value| serde_json::from_value::(value.clone())) + .transpose() + .map_err(|error| RunAdmissionError::Persistence(error.to_string()))?; + crate::session_fork_coordinator::authorize_provider_fork_session( + &mut *tx, + &record.user_id, + &record.session_id, + provider_owner.as_ref(), + self.scope_encryptor.as_deref(), + crate::session_fork_coordinator::ProviderSessionScopeRequirement::ExistingConversation, + ) + .await + .map_err(|error| RunAdmissionError::Persistence(error.to_string()))?; let mut collaborator_event_rows = vec![]; let mut collaborator_events_on_new_run = false; let mut stage_receipt = None; diff --git a/crates/services/src/runtime_maintenance.rs b/crates/services/src/runtime_maintenance.rs index 4506211fed..47ad8c1df6 100644 --- a/crates/services/src/runtime_maintenance.rs +++ b/crates/services/src/runtime_maintenance.rs @@ -3,7 +3,8 @@ //! //! This worker is deliberately limited to explicit user delete intents, //! expired diagnostics, and runtime objects whose owner was already removed. -//! It does not transition or end agent_sessions, and it does not expire +//! It also cancels expired, unused provider fork preparations through their +//! existing abort and Session deletion owners. It does not expire //! session traces. //! Durable conversation history remains available until the user deletes the //! session through the normal API. @@ -46,6 +47,8 @@ pub struct RuntimeMaintenanceSweepResult { pub observation_collision_receipts_expired: u64, /// Expired fork-retention pins released after their grace period. pub expired_fork_pins_released: u64, + /// Unactivated provider forks cancelled after their retry window. + pub expired_provider_prepares_aborted: u64, /// Fork records whose child session no longer exists. pub orphaned_forks_collected: u64, /// Unpinned manifest nodes whose owner session no longer exists. @@ -131,6 +134,13 @@ pub async fn maintain_runtime_storage( expire_session_attachments(pool, policy.batch_limit).await, ); if let Some(coordinator) = fork_coordinator { + result.expired_provider_prepares_aborted = record_runtime_storage_maintenance( + &mut result.cleanup_errors, + "abort_expired_provider_prepares", + coordinator + .abort_expired_provider_prepares(policy.batch_limit) + .await, + ); result.expired_fork_pins_released = record_runtime_storage_maintenance( &mut result.cleanup_errors, "release_expired_fork_pins", @@ -163,6 +173,7 @@ impl RuntimeMaintenanceSweepResult { .saturating_add(self.model_request_context_events_expired) .saturating_add(self.observation_collision_receipts_expired) .saturating_add(self.expired_fork_pins_released) + .saturating_add(self.expired_provider_prepares_aborted) .saturating_add(self.orphaned_forks_collected) .saturating_add(self.orphaned_manifests_collected) .saturating_add(self.work_branch_control_operations_expired) diff --git a/crates/services/src/session_fork_coordinator.rs b/crates/services/src/session_fork_coordinator.rs index 3af1b3c79a..1029c8cac5 100644 --- a/crates/services/src/session_fork_coordinator.rs +++ b/crates/services/src/session_fork_coordinator.rs @@ -29,6 +29,7 @@ const MAX_REASON_BYTES: usize = 1_024; const MAX_LIST_LIMIT: u32 = 200; const MAX_WRITER_TTL: Duration = Duration::from_secs(15 * 60); const MAX_ABORT_GRACE: Duration = Duration::from_secs(30 * 24 * 60 * 60); +const PROVIDER_PREPARE_TTL_SECONDS: i64 = 24 * 60 * 60; #[derive(Debug, Error)] pub enum SessionForkCoordinatorError { @@ -36,6 +37,21 @@ pub enum SessionForkCoordinatorError { Invalid(String), #[error("session fork was not found")] NotFound, + #[error("session fork is outside the authorized provider scope")] + Forbidden, + #[error("session creation quota exceeded: {reason}")] + QuotaDenied { + limit: crate::resource_governor::ResourceLimitKind, + reason: String, + }, + #[error("session fork was aborted before activation")] + Aborted, + #[error("session fork cursor projection is not ready")] + CursorNotReady, + #[error("session fork cursor is no longer available")] + CursorUnavailable, + #[error("run does not support a canonical fork cursor")] + CursorUnsupported, #[error("session fork conflicts with existing child state")] Conflict, #[error("session fork idempotency key was reused for a different request")] @@ -70,22 +86,261 @@ pub struct PrepareSessionForkV1 { pub reason: String, } +/// Activated context without a writer lease or execution authority. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct SessionForkContinuationV1 { + pub manifest: SessionForkManifestV1, + pub child_head: astra_turn_types::SessionContextHeadV1, +} + #[derive(Clone)] pub struct DatabaseSessionForkCoordinator { pool: SharedPool, context: Arc, + scope_encryptor: Option>, } impl DatabaseSessionForkCoordinator { pub fn new(pool: SharedPool, context: Arc) -> Self { - Self { pool, context } + Self { + pool, + context, + scope_encryptor: None, + } + } + + pub fn with_scope_encryptor( + mut self, + encryptor: Arc, + ) -> Self { + self.scope_encryptor = Some(encryptor); + self + } + + /// Resolve an exact root-run commit, never the Session's current head. + pub async fn resolve_run_cursor( + &self, + key: &SessionKeyV1, + run_id: &str, + provider_scope: Option<&crate::runs::ProviderRunOwner>, + ) -> Result { + validate_key_and_id(key, run_id)?; + let run = sqlx::query( + "SELECT parent_run_id, root_run_id, depth, status FROM agent_runs + WHERE user_id = ? AND session_id = ? AND run_id = ?", + ) + .bind(&key.owner_user_id) + .bind(&key.session_id) + .bind(run_id) + .fetch_optional(self.pool.get()) + .await + .map_err(|source| database_error("load_fork_source_run", source))? + .ok_or(SessionForkCoordinatorError::NotFound)?; + if let Some(required) = provider_scope { + authorize_provider_fork_session( + self.pool.get(), + &key.owner_user_id, + &key.session_id, + Some(required), + self.scope_encryptor.as_deref(), + ProviderSessionScopeRequirement::ForkSource, + ) + .await?; + } + let parent: Option = run + .try_get("parent_run_id") + .map_err(|source| database_error("decode_fork_source_parent", source))?; + let root: String = run + .try_get("root_run_id") + .map_err(|source| database_error("decode_fork_source_root", source))?; + let depth: i32 = run + .try_get("depth") + .map_err(|source| database_error("decode_fork_source_depth", source))?; + let status: String = run + .try_get("status") + .map_err(|source| database_error("decode_fork_source_status", source))?; + if parent.is_some() || root != run_id || depth != 0 { + return Err(SessionForkCoordinatorError::CursorUnsupported); + } + if status != "completed" { + return Err(if matches!(status.as_str(), "failed" | "cancelled") { + SessionForkCoordinatorError::CursorUnsupported + } else { + SessionForkCoordinatorError::CursorNotReady + }); + } + let rows = sqlx::query( + "SELECT DISTINCT canonical_completed_turn, canonical_conversation_seq, canonical_root_hash + FROM session_transcript_items WHERE user_id = ? AND session_id = ? AND run_id = ? LIMIT 2", + ) + .bind(&key.owner_user_id) + .bind(&key.session_id) + .bind(run_id) + .fetch_all(self.pool.get()) + .await + .map_err(|source| database_error("load_fork_source_commit", source))?; + if rows.len() != 1 { + return Err( + if rows.is_empty() && self.context.load_head(key).await?.is_none() { + SessionForkCoordinatorError::CursorUnsupported + } else { + SessionForkCoordinatorError::CursorNotReady + }, + ); + } + let row = &rows[0]; + let turn: Option = row + .try_get("canonical_completed_turn") + .map_err(|source| database_error("decode_fork_source_turn", source))?; + let sequence: Option = row + .try_get("canonical_conversation_seq") + .map_err(|source| database_error("decode_fork_source_sequence", source))?; + let root: Option = row + .try_get("canonical_root_hash") + .map_err(|source| database_error("decode_fork_source_commit_root", source))?; + let (Some(turn), Some(sequence), Some(root)) = (turn, sequence, root) else { + return Err(if self.context.load_head(key).await?.is_none() { + SessionForkCoordinatorError::CursorUnsupported + } else { + SessionForkCoordinatorError::CursorNotReady + }); + }; + let row = sqlx::query( + "SELECT manifest_json FROM conversation_manifest_nodes FORCE INDEX (uq_context_manifest_root) + WHERE isolation_domain = ? AND owner_user_id = ? AND session_id = ? + AND branch_id = ? AND manifest_root = ? AND reachable = 1", + ) + .bind(&key.isolation_domain) + .bind(&key.owner_user_id) + .bind(&key.session_id) + .bind(&key.branch_id) + .bind(&root) + .fetch_optional(self.pool.get()) + .await + .map_err(|source| database_error("load_fork_source_manifest", source))? + .ok_or(SessionForkCoordinatorError::CursorUnavailable)?; + let node: astra_turn_types::ContextManifestNodeV1 = + decode_json_row(&row, "manifest_json", "fork_source_manifest")?; + node.validate() + .map_err(|error| SessionForkCoordinatorError::NeedsRepair(error.to_string()))?; + let cursor = node.cursor(); + if node.key != *key + || cursor.canonical_root_hash != root + || i64::from(cursor.completed_turn) != turn + || i64::try_from(cursor.conversation_seq).ok() != Some(sequence) + { + return Err(SessionForkCoordinatorError::CursorUnavailable); + } + Ok(cursor) + } + + /// The authenticated Session binding owns provider scope. Preparing the + /// fork separately validates this exact canonical cursor under its write + /// fence, so authorization does not depend on a delayed Run projection. + pub async fn authorize_provider_cursor( + &self, + key: &SessionKeyV1, + cursor: &SessionCursorV1, + required: &crate::runs::ProviderRunOwner, + ) -> Result<(), SessionForkCoordinatorError> { + key.validate() + .map_err(|error| SessionForkCoordinatorError::Invalid(error.to_string()))?; + if !key.validates_cursor(cursor) { + return Err(SessionForkCoordinatorError::Forbidden); + } + authorize_provider_fork_session( + self.pool.get(), + &key.owner_user_id, + &key.session_id, + Some(required), + self.scope_encryptor.as_deref(), + ProviderSessionScopeRequirement::ForkSource, + ) + .await + } + + /// Activate immutable context only. Ordinary chat admits fresh Run-owned Context authority. + pub async fn activate_for_continuation( + &self, + parent_key: &SessionKeyV1, + fork_id: &str, + ) -> Result { + let stored = self.load(parent_key, fork_id).await?; + self.activate_stored_fork(stored).await + } + + async fn activate_stored_fork( + &self, + stored: SessionForkManifestV1, + ) -> Result { + if stored.state == SessionForkStateV1::Aborted { + return Err(SessionForkCoordinatorError::Aborted); + } + let child_head = if stored.state == SessionForkStateV1::Prepared { + match self.context.activate_fork(&stored).await { + Ok(head) => head, + Err(error) => { + if self.load(&stored.parent_key, &stored.fork_id).await?.state + == SessionForkStateV1::Aborted + { + return Err(SessionForkCoordinatorError::Aborted); + } + return Err(error.into()); + } + } + } else { + self.context + .load_head(&stored.child_key) + .await? + .ok_or_else(|| { + SessionForkCoordinatorError::NeedsRepair( + "active fork has no child context head".into(), + ) + })? + }; + if child_head.key != stored.child_key + || !stored.child_key.validates_cursor(&child_head.cursor) + { + return Err(SessionForkCoordinatorError::NeedsRepair( + "active fork child head escaped its owner".into(), + )); + } + let manifest = self.load(&stored.parent_key, &stored.fork_id).await?; + if manifest.state != SessionForkStateV1::Active { + return Err(SessionForkCoordinatorError::NeedsRepair( + "fork activation did not durably transition to active".into(), + )); + } + Ok(SessionForkContinuationV1 { + manifest, + child_head, + }) } pub async fn prepare( &self, request: &PrepareSessionForkV1, + ) -> Result { + self.prepare_admitted(request, None, false).await + } + + /// HTTP creation supplies only authenticated scope and uses the same + /// durable quota reservation as ordinary Session creation. + pub async fn prepare_admitted( + &self, + request: &PrepareSessionForkV1, + provider_owner: Option<&crate::runs::ProviderRunOwner>, + enforce_creation_quota: bool, ) -> Result { validate_prepare_request(request)?; + if let Some(owner) = provider_owner { + self.authorize_provider_cursor( + &request.parent_key, + &request.expected_parent_cursor, + owner, + ) + .await?; + } let request_hash = stable_hash(b"astra.prepare-session-fork.v1\0", request)?; let idempotency_hash = identity_hash("prepare", &request.idempotency_key); let mut connection = CancellationSafePoolConnection::acquire(self.pool.get()) @@ -95,11 +350,38 @@ impl DatabaseSessionForkCoordinator { .begin() .await .map_err(|source| database_error("begin_prepare_fork", source))?; - crate::session_context_coordinator::admit_session_lifecycle_writes( + if let Err(error) = crate::session_context_coordinator::admit_session_lifecycle_writes( &mut tx, [&request.parent_key, &request.child_key], ) - .await?; + .await + { + if let Some(row) = sqlx::query( + "SELECT request_hash, manifest_json FROM session_forks + WHERE isolation_domain = ? AND owner_user_id = ? + AND parent_session_id = ? AND parent_branch_id = ? AND idempotency_hash = ?", + ) + .bind(&request.parent_key.isolation_domain) + .bind(&request.parent_key.owner_user_id) + .bind(&request.parent_key.session_id) + .bind(&request.parent_key.branch_id) + .bind(&idempotency_hash) + .fetch_optional(&mut *tx) + .await + .map_err(|source| database_error("load_aborted_prepare_retry", source))? + { + verify_request_hash(&row, &request_hash)?; + let manifest: SessionForkManifestV1 = + decode_json_row(&row, "manifest_json", "aborted_prepare_retry")?; + validate_stored_manifest(&manifest, &request.parent_key)?; + if manifest.state == SessionForkStateV1::Aborted + && manifest.provider_owner.as_ref() == provider_owner + { + return Err(SessionForkCoordinatorError::Aborted); + } + } + return Err(error.into()); + } if let Some(row) = sqlx::query( "SELECT request_hash, manifest_json FROM session_forks @@ -119,6 +401,12 @@ impl DatabaseSessionForkCoordinator { verify_request_hash(&row, &request_hash)?; let manifest = decode_json_row(&row, "manifest_json", "fork_manifest")?; validate_stored_manifest(&manifest, &request.parent_key)?; + if manifest.provider_owner.as_ref() != provider_owner { + return Err(SessionForkCoordinatorError::Forbidden); + } + if manifest.state == SessionForkStateV1::Aborted { + return Err(SessionForkCoordinatorError::Aborted); + } tx.commit() .await .map_err(|source| database_error("commit_prepare_retry", source))?; @@ -128,6 +416,18 @@ impl DatabaseSessionForkCoordinator { let parent_head = lock_parent_head(&mut tx, request).await?; ensure_child_is_empty(&mut tx, &request.child_key).await?; + if enforce_creation_quota { + let quota = + crate::resource_governor::DatabaseResourceGovernor::reserve_session_creation( + &mut tx, + &request.parent_key.owner_user_id, + ) + .await + .map_err(|source| database_error("reserve_fork_creation_quota", source))?; + if let crate::resource_governor::LimitCheck::Denied { limit, reason } = quota { + return Err(SessionForkCoordinatorError::QuotaDenied { limit, reason }); + } + } let now = database_now_ms(&mut tx).await?; let manifest = SessionForkManifestV1 { schema_version: SESSION_FORK_MANIFEST_SCHEMA_VERSION, @@ -135,6 +435,7 @@ impl DatabaseSessionForkCoordinator { parent_key: request.parent_key.clone(), child_key: request.child_key.clone(), parent_head, + provider_owner: provider_owner.cloned(), dimensions: request.dimensions.clone(), excluded_authority: vec![ ForkExcludedAuthorityV1::Run, @@ -176,8 +477,8 @@ impl DatabaseSessionForkCoordinator { sqlx::query( "INSERT INTO conversation_manifest_pins (isolation_domain, owner_user_id, pin_id, parent_session_id, - parent_branch_id, manifest_root, pin_state) - VALUES (?, ?, ?, ?, ?, ?, 'prepared')", + parent_branch_id, manifest_root, pin_state, grace_expires_at_ms) + VALUES (?, ?, ?, ?, ?, ?, 'prepared', ?)", ) .bind(&manifest.parent_key.isolation_domain) .bind(&manifest.parent_key.owner_user_id) @@ -185,6 +486,7 @@ impl DatabaseSessionForkCoordinator { .bind(&manifest.parent_key.session_id) .bind(&manifest.parent_key.branch_id) .bind(&manifest.parent_head.latest_manifest_root) + .bind(provider_owner.map(|_| now + PROVIDER_PREPARE_TTL_SECONDS * 1000)) .execute(&mut *tx) .await .map_err(|source| database_error("pin_prepared_fork", source))?; @@ -209,21 +511,9 @@ impl DatabaseSessionForkCoordinator { actor .validate_for(&stored.child_key) .map_err(|error| SessionForkCoordinatorError::Invalid(error.to_string()))?; - if stored.state == SessionForkStateV1::Aborted { - return Err(SessionForkCoordinatorError::Conflict); - } - let child_head = if stored.state == SessionForkStateV1::Prepared { - self.context.activate_fork(&stored).await? - } else { - self.context - .load_head(&stored.child_key) - .await? - .ok_or_else(|| { - SessionForkCoordinatorError::NeedsRepair( - "active fork has no child context head".into(), - ) - })? - }; + let continuation = self.activate_stored_fork(stored).await?; + let stored = continuation.manifest; + let child_head = continuation.child_head; let writer_idempotency = identity_hash( "activate-writer", &format!("{}\0{}", stored.fork_id, actor.actor_id), @@ -302,11 +592,25 @@ impl DatabaseSessionForkCoordinator { child_session_id, child_branch_id, ); - crate::session_context_coordinator::admit_session_lifecycle_writes( - &mut tx, - [parent_key, &child_key], - ) - .await?; + // Cancellation needs the same lock order as activation, but can replay + // after its own child deletion fence and after parent deletion. + let mut keys = [parent_key, &child_key]; + keys.sort_unstable_by_key(|key| { + (&key.owner_user_id, &key.session_id, &key.isolation_domain) + }); + let mut child_writable = false; + for key in keys { + let state = crate::storage::lock_agent_session_write_fence_state( + &mut tx, + &key.session_id, + &key.owner_user_id, + ) + .await + .map_err(|source| database_error("lock_abort_lifecycle_fence", source))?; + if key == &child_key { + child_writable = state == crate::storage::AgentSessionWriteFenceState::Writable; + } + } let row = lock_fork(&mut tx, parent_key, fork_id).await?; let mut manifest: SessionForkManifestV1 = decode_json_row(&row, "manifest_json", "fork_manifest")?; @@ -321,6 +625,31 @@ impl DatabaseSessionForkCoordinator { if manifest.state != SessionForkStateV1::Prepared { return Err(SessionForkCoordinatorError::Conflict); } + if !child_writable { + return Err(SessionForkCoordinatorError::Conflict); + } + if manifest.provider_owner.is_some() { + let used: bool = sqlx::query_scalar( + "SELECT (EXISTS (SELECT 1 FROM agent_runs WHERE user_id = ? AND session_id = ?)) + OR (EXISTS (SELECT 1 FROM session_context_heads + WHERE isolation_domain = ? AND owner_user_id = ? AND session_id = ? AND branch_id = ? + AND (head_json IS NOT NULL OR fork_base_json IS NOT NULL + OR active_writer_json IS NOT NULL OR active_reservation_json IS NOT NULL)))", + ).bind(&child_key.owner_user_id).bind(&child_key.session_id) + .bind(&child_key.isolation_domain).bind(&child_key.owner_user_id) + .bind(&child_key.session_id).bind(&child_key.branch_id) + .fetch_one(&mut *tx).await.map_err(|source| database_error("check_cancelled_child_unused", source))?; + if used { + return Err(SessionForkCoordinatorError::Conflict); + } + crate::session_lifecycle::mark_session_deleting_tx( + &mut tx, + &child_key.session_id, + &child_key.owner_user_id, + ) + .await + .map_err(SessionForkCoordinatorError::NeedsRepair)?; + } let now = database_now_ms(&mut tx).await?; let grace_expires_at_ms = checked_expiry(now, grace)?; manifest.state = SessionForkStateV1::Aborted; @@ -459,6 +788,56 @@ impl DatabaseSessionForkCoordinator { Ok(result.rows_affected()) } + /// A lost provider prepare has a bounded retry window. Reuse the ordinary + /// Prepared -> Aborted transition and durable Session delete intent. + pub async fn abort_expired_provider_prepares( + &self, + limit: u32, + ) -> Result { + if limit == 0 || limit > 10_000 { + return Err(SessionForkCoordinatorError::Invalid( + "cleanup limit must be between 1 and 10000".into(), + )); + } + let rows = sqlx::query( + "SELECT forks.manifest_json FROM conversation_manifest_pins pins + FORCE INDEX (idx_manifest_pins_grace) + JOIN session_forks forks ON forks.isolation_domain = pins.isolation_domain + AND forks.owner_user_id = pins.owner_user_id AND forks.fork_id = pins.pin_id + WHERE pins.pin_state = 'prepared' + AND pins.grace_expires_at_ms <= CAST(UNIX_TIMESTAMP(NOW(6)) * 1000 AS SIGNED) + ORDER BY pins.grace_expires_at_ms ASC LIMIT ?", + ) + .bind(i64::from(limit)) + .fetch_all(self.pool.get()) + .await + .map_err(|source| database_error("load_expired_provider_prepares", source))?; + let mut aborted = 0; + for row in rows { + let manifest: SessionForkManifestV1 = + decode_json_row(&row, "manifest_json", "expired_provider_prepare")?; + if manifest.provider_owner.is_none() { + continue; + } + match self + .abort( + &manifest.parent_key, + &manifest.fork_id, + Duration::from_secs(1), + "provider prepare expired", + ) + .await + { + Ok(_) => aborted += 1, + Err( + SessionForkCoordinatorError::Conflict | SessionForkCoordinatorError::NotFound, + ) => {} + Err(error) => return Err(error), + } + } + Ok(aborted) + } + /// Collect manifest records whose owning session was deleted after every /// fork retention pin was released. Node-local payloads are deleted with /// their manifest row, without JSON scans or per-segment refcount fanout. @@ -553,6 +932,13 @@ impl DatabaseSessionForkCoordinator { WHERE agent_sessions.user_id = session_forks.owner_user_id AND agent_sessions.session_id = session_forks.child_session_id ) + AND (session_forks.state <> 'aborted' + OR JSON_EXTRACT(session_forks.manifest_json, '$.provider_owner') IS NULL + OR NOT EXISTS ( + SELECT 1 FROM agent_sessions parent + WHERE parent.user_id = session_forks.owner_user_id + AND parent.session_id = session_forks.parent_session_id + )) AND NOT EXISTS ( SELECT 1 FROM session_forks descendants WHERE descendants.isolation_domain = @@ -677,6 +1063,142 @@ fn validate_prepare_request( Ok(()) } +pub fn authorize_manifest_provider( + manifest: &SessionForkManifestV1, + required: &crate::runs::ProviderRunOwner, +) -> Result<(), SessionForkCoordinatorError> { + if manifest.provider_owner.as_ref() == Some(required) { + Ok(()) + } else { + Err(SessionForkCoordinatorError::Forbidden) + } +} + +fn provider_session_namespace(user_id: &str) -> bool { + user_id.starts_with("provider_authorized:") +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ProviderSessionScopeRequirement { + ExistingConversation, + ForkSource, +} + +/// Common Run admission, including replay. Forks use their retained manifest; +/// new provider Sessions use their server-authenticated scope proof. Existing +/// unbound ordinary conversations retain their previous chat contract, but +/// cannot become a source for a new provider fork. +pub async fn authorize_provider_fork_session<'e, E>( + executor: E, + user_id: &str, + session_id: &str, + required: Option<&crate::runs::ProviderRunOwner>, + encryptor: Option<&crate::auth::FernetTokenEncryptor>, + requirement: ProviderSessionScopeRequirement, +) -> Result<(), SessionForkCoordinatorError> +where + E: sqlx::Executor<'e, Database = MySql>, +{ + // Legacy user tokens have a distinct owner namespace and retain the + // existing subject-owned Session contract. MOI's request-signed Sessions + // cannot be reached through it, and new provider fork control requires + // request-bound authentication rather than legacy credentials. + if user_id.starts_with("external_authorized:moi:") { + return Ok(()); + } + // Local identities cannot address request-signed provider Sessions. Their + // ordinary admission already validates Session ownership and lifecycle. + if required.is_none() && !provider_session_namespace(user_id) { + return Ok(()); + } + let key = SessionKeyV1::owner_session("server", user_id, session_id, "main"); + let row = sqlx::query( + "SELECT JSON_EXTRACT(s.`metadata`, '$.__astra_provider_scope_proof_v1') IS NOT NULL AS has_scope_binding, + JSON_EXTRACT(s.`metadata`, '$.fork_id') IS NOT NULL AS has_fork_metadata, + CASE WHEN OCTET_LENGTH(JSON_UNQUOTE(JSON_EXTRACT(s.`metadata`, '$.__astra_provider_scope_proof_v1'))) = 43 + THEN JSON_UNQUOTE(JSON_EXTRACT(s.`metadata`, '$.__astra_provider_scope_proof_v1')) + ELSE NULL END AS scope_proof, + CAST(f.manifest_json AS CHAR) AS scope_manifest + FROM agent_sessions s LEFT JOIN session_forks f + ON f.isolation_domain = ? AND f.owner_user_id = s.user_id + AND f.child_session_id = s.session_id AND f.child_branch_id = ? + WHERE s.user_id = ? AND s.session_id = ?", + ) + .bind(&key.isolation_domain) + .bind(&key.branch_id) + .bind(user_id) + .bind(session_id) + .fetch_optional(executor) + .await + .map_err(|source| database_error("load_provider_session_scope", source))?; + let Some(row) = row else { + return if required.is_some() || provider_session_namespace(user_id) { + Err(SessionForkCoordinatorError::Forbidden) + } else { + Ok(()) + }; + }; + let raw_manifest: Option = row + .try_get("scope_manifest") + .map_err(|source| database_error("decode_provider_session_scope", source))?; + let Some(raw_manifest) = raw_manifest else { + let Some(owner) = required else { + return if provider_session_namespace(user_id) { + Err(SessionForkCoordinatorError::Forbidden) + } else { + Ok(()) + }; + }; + let has_scope_binding: bool = row + .try_get("has_scope_binding") + .map_err(|source| database_error("decode_provider_session_scope", source))?; + let has_fork_metadata: bool = row + .try_get("has_fork_metadata") + .map_err(|source| database_error("decode_provider_session_scope", source))?; + if !has_scope_binding + && !has_fork_metadata + && requirement == ProviderSessionScopeRequirement::ExistingConversation + { + return Ok(()); + } + let proof: Option = row + .try_get("scope_proof") + .map_err(|source| database_error("decode_provider_session_scope", source))?; + return if proof + .as_deref() + .zip(encryptor) + .is_some_and(|(proof, encryptor)| { + encryptor.verify_provider_session_scope( + user_id, + session_id, + &owner.provider_id, + &owner.provider_scope_id, + proof, + ) + }) { + Ok(()) + } else { + Err(SessionForkCoordinatorError::Forbidden) + }; + }; + let manifest: SessionForkManifestV1 = serde_json::from_str(&raw_manifest) + .map_err(|error| SessionForkCoordinatorError::NeedsRepair(error.to_string()))?; + validate_stored_manifest(&manifest, &manifest.parent_key)?; + if manifest.child_key != key { + return Err(SessionForkCoordinatorError::Forbidden); + } + if let Some(owner) = manifest.provider_owner.as_ref() { + if required != Some(owner) || manifest.state != SessionForkStateV1::Active { + return Err(SessionForkCoordinatorError::Forbidden); + } + } else if required.is_some() || provider_session_namespace(user_id) { + // A provider fork without retained scope cannot be claimed by the + // first request, regardless of which provider credential was used. + return Err(SessionForkCoordinatorError::Forbidden); + } + Ok(()) +} + async fn insert_fork_child_session( tx: &mut Transaction<'_, MySql>, manifest: &SessionForkManifestV1, @@ -1243,6 +1765,18 @@ mod tests { .await .expect("activate fork"); assert_eq!(activation.manifest.state, SessionForkStateV1::Active); + let copied_transcript: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM session_transcript_items WHERE user_id = ? AND session_id = ?", + ) + .bind(&owner) + .bind(&child_key.session_id) + .fetch_one(pool.get()) + .await + .expect("read child transcript"); + assert_eq!( + copied_transcript, 0, + "baseline must not copy parent transcript rows" + ); let child_manifest_count_after: i64 = sqlx::query_scalar( "SELECT COUNT(*) FROM conversation_manifest_nodes WHERE isolation_domain = ? AND owner_user_id = ? diff --git a/crates/services/src/session_lifecycle.rs b/crates/services/src/session_lifecycle.rs index ab743d49be..dc9c51f4fd 100644 --- a/crates/services/src/session_lifecycle.rs +++ b/crates/services/src/session_lifecycle.rs @@ -90,6 +90,11 @@ const SESSION_DELETE_CHILD_FORKS_SELECT_SQL: &str = "SELECT isolation_domain, fork_id, parent_session_id FROM session_forks WHERE child_session_id = ? AND owner_user_id = ? + AND (state <> 'aborted' OR JSON_EXTRACT(manifest_json, '$.provider_owner') IS NULL OR NOT EXISTS ( + SELECT 1 FROM agent_sessions parent + WHERE parent.user_id = session_forks.owner_user_id + AND parent.session_id = session_forks.parent_session_id + )) AND NOT EXISTS ( SELECT 1 FROM session_forks descendants WHERE descendants.isolation_domain = session_forks.isolation_domain @@ -1025,12 +1030,26 @@ async fn mark_session_deleting( .await .map_err(|source| format!("delete_session.mark_deleting.lock_fence: {source}"))?; - ensure_no_live_work_recovery_points(&mut tx, session_id, user_id).await?; + mark_session_deleting_tx(&mut tx, session_id, user_id).await?; + tx.commit() + .await + .map_err(|source| format!("delete_session.mark_deleting.commit: {source}"))?; + Ok(()) +} + +/// Persist a delete intent under the caller's already-held Session fence. +/// Prepared fork cancellation uses this same owner; maintenance completes it. +pub(crate) async fn mark_session_deleting_tx( + tx: &mut sqlx::Transaction<'_, MySql>, + session_id: &str, + user_id: &str, +) -> Result<(), String> { + ensure_no_live_work_recovery_points(tx, session_id, user_id).await?; let result = query(MARK_SESSION_DELETING_SQL) .bind(session_id) .bind(user_id) - .execute(&mut *tx) + .execute(&mut **tx) .await .map_err(|source| format!("delete_session.mark_deleting: {source}"))?; @@ -1039,7 +1058,7 @@ async fn mark_session_deleting( query("SELECT 1 FROM agent_sessions WHERE session_id = ? AND user_id = ? LIMIT 1") .bind(session_id) .bind(user_id) - .fetch_optional(&mut *tx) + .fetch_optional(&mut **tx) .await .map_err(|source| format!("delete_session.mark_deleting.confirm: {source}"))? .is_some(); @@ -1053,14 +1072,10 @@ async fn mark_session_deleting( query(MARK_SESSION_DELETE_FENCE_SQL) .bind(session_id) .bind(user_id) - .execute(&mut *tx) + .execute(&mut **tx) .await .map_err(|source| format!("delete_session.mark_deleting.persist_fence: {source}"))?; - tx.commit() - .await - .map_err(|source| format!("delete_session.mark_deleting.commit: {source}"))?; - Ok(()) } diff --git a/crates/services/tests/provider_session_lookup_db_it.rs b/crates/services/tests/provider_session_lookup_db_it.rs index 06ba8ef993..22a5b56d50 100644 --- a/crates/services/tests/provider_session_lookup_db_it.rs +++ b/crates/services/tests/provider_session_lookup_db_it.rs @@ -3,6 +3,7 @@ mod common; use astra_services::{DatabaseSessionService, SessionCreateRequestData, SessionService}; use axum::http::StatusCode; use serial_test::serial; +use std::sync::Arc; use uuid::Uuid; #[tokio::test] @@ -15,8 +16,19 @@ async fn provider_session_creation_is_idempotent_across_service_instances() { ProviderSessionCreationIdentity, }; let (pool, settings) = common::setup_pool_and_settings().await; - let first = DatabaseSessionService::new(settings.clone()).with_pool(pool.clone()); - let second = DatabaseSessionService::new(settings).with_pool(pool.clone()); + astra_services::resource_governor::DatabaseResourceGovernor::new(pool.clone()) + .ensure_tables() + .await + .unwrap(); + let scope_key = Arc::new( + astra_services::FernetTokenEncryptor::new("provider-session-scope-test-key").unwrap(), + ); + let first = DatabaseSessionService::new(settings.clone()) + .with_pool(pool.clone()) + .with_scope_encryptor(scope_key.clone()); + let second = DatabaseSessionService::new(settings) + .with_pool(pool.clone()) + .with_scope_encryptor(scope_key); let user_id = format!("creation-{}", Uuid::new_v4().simple()); let principal = AuthPrincipal { user: AuthUserRecord { @@ -54,6 +66,47 @@ async fn provider_session_creation_is_idempotent_across_service_instances() { assert_ne!(a.created, b.created, "exactly one call creates the session"); assert_eq!(a.admission_facts.active_plan_id, None); assert_eq!(b.admission_facts.active_plan_id, None); + let proof_key = "__astra_provider_scope_proof_v1"; + let proof = a + .session + .metadata + .get(proof_key) + .cloned() + .expect("server scope proof"); + let replacement = first + .update_session( + a.session.session_id.clone(), + user_id.clone(), + astra_services::SessionUpdateRequestData { + title: None, + metadata: Some(serde_json::Map::from_iter([( + "user_key".into(), + serde_json::json!("value"), + )])), + metadata_patch: None, + status: None, + }, + ) + .await + .expect("ordinary metadata replacement"); + assert_eq!(replacement.metadata.get(proof_key), Some(&proof)); + let forged = first + .update_session( + a.session.session_id.clone(), + user_id.clone(), + astra_services::SessionUpdateRequestData { + title: None, + metadata: None, + metadata_patch: Some(serde_json::Map::from_iter([( + proof_key.into(), + serde_json::Value::Null, + )])), + status: None, + }, + ) + .await + .expect_err("client must not remove server scope proof"); + assert_eq!(forged.0, StatusCode::BAD_REQUEST); let active_plan_id = Uuid::new_v4().to_string(); sqlx::query( "UPDATE agent_sessions SET active_plan_id = ? WHERE user_id = ? AND session_id = ?", @@ -93,9 +146,14 @@ async fn provider_session_creation_is_idempotent_across_service_instances() { StatusCode::TOO_MANY_REQUESTS ); first - .delete_session(a.session.session_id, user_id) + .delete_session(a.session.session_id, user_id.clone()) .await .expect("delete test session"); + sqlx::query("DELETE FROM resource_usage WHERE user_id = ?") + .bind(&user_id) + .execute(pool.get()) + .await + .unwrap(); } #[tokio::test] diff --git a/crates/services/tests/session_creation_quota_db_it.rs b/crates/services/tests/session_creation_quota_db_it.rs new file mode 100644 index 0000000000..104c3ab350 --- /dev/null +++ b/crates/services/tests/session_creation_quota_db_it.rs @@ -0,0 +1,294 @@ +use astra_core::{MatrixOneSettings, SharedPool}; +use astra_services::{ + AuthPrincipal, AuthPrincipalOrigin, AuthProviderAuthorizedRequestContext, AuthUserRecord, + DatabaseSessionService, ProviderSessionCreationIdentity, SessionCreateRequestData, + SessionService, + resource_governor::{DatabaseResourceGovernor, LimitCheck, ResourceGovernor, ResourceLimits}, +}; +use axum::http::StatusCode; +use sqlx::{Connection, MySqlConnection}; +use std::sync::Arc; +use uuid::Uuid; + +#[tokio::test] +#[ignore = "requires live DB: run with ASTRA_TEST_DB_IT=1"] +async fn provider_creation_reserves_daily_quota_once_in_its_transaction() { + assert_eq!(std::env::var("ASTRA_TEST_DB_IT").as_deref(), Ok("1")); + let mut settings = MatrixOneSettings::from_env(); + settings.database = "mysql".into(); + let mut admin = MySqlConnection::connect(&settings.database_url_with_password()) + .await + .unwrap(); + let database = format!("astra_test_creation_quota_{}", Uuid::new_v4().simple()); + sqlx::query(sqlx::AssertSqlSafe(format!("CREATE DATABASE `{database}`"))) + .execute(&mut admin) + .await + .unwrap(); + settings.database = database.clone(); + astra_services::storage::ensure_core_schema(&settings, "mysql") + .await + .unwrap(); + let pool = SharedPool::new(&settings).await.unwrap(); + let governor = DatabaseResourceGovernor::new(pool.clone()); + let sessions = DatabaseSessionService::new(settings) + .with_pool(pool.clone()) + .with_scope_encryptor(Arc::new( + astra_services::FernetTokenEncryptor::new("provider-session-scope-test-key").unwrap(), + )); + // Internal core-only users retain the unmetered creation API. HTTP quota + // admission instead fails closed if its required resource tables are absent. + let internal = sessions + .create_session( + "quota-internal".into(), + SessionCreateRequestData { + agent_id: None, + title: None, + metadata: None, + }, + ) + .await; + let unavailable = sessions + .create_session_with_resource_quota( + "quota-unavailable".into(), + SessionCreateRequestData { + agent_id: None, + title: None, + metadata: None, + }, + &governor, + ) + .await; + let unavailable_sessions: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM agent_sessions WHERE user_id = ?") + .bind("quota-unavailable") + .fetch_one(pool.get()) + .await + .unwrap(); + governor.ensure_tables().await.unwrap(); + let user_id = format!("quota-{}", Uuid::new_v4().simple()); + let principal = AuthPrincipal { + user: AuthUserRecord { + user_id: user_id.clone(), + username: user_id.clone(), + email: "quota@example.test".into(), + display_name: None, + }, + session_id: None, + execution_continuation: None, + origin: AuthPrincipalOrigin::ProviderAuthorizedRequest( + AuthProviderAuthorizedRequestContext { + provider_id: "provider-a".into(), + external_subject: user_id.clone(), + provider_scope_id: "scope-a".into(), + request_authorization_id: "quota-request".into(), + edge_agent_id: None, + }, + ), + }; + governor + .set_limits( + &user_id, + ResourceLimits { + max_sessions_per_day: 1, + ..Default::default() + }, + ) + .await; + let request = SessionCreateRequestData { + agent_id: None, + title: None, + metadata: None, + }; + let first_identity = + ProviderSessionCreationIdentity::from_principal(&principal, "first").unwrap(); + let first = sessions + .create_provider_session(first_identity.clone(), request.clone(), LimitCheck::Allowed) + .await; + let proof_key = "__astra_provider_scope_proof_v1"; + let original_proof = first + .as_ref() + .ok() + .and_then(|result| result.session.metadata.get(proof_key).cloned()); + let (replacement, forged) = if let Ok(result) = first.as_ref() { + let session_id = result.session.session_id.clone(); + let replacement = sessions + .update_session( + session_id.clone(), + user_id.clone(), + astra_services::SessionUpdateRequestData { + title: None, + metadata: Some(serde_json::Map::from_iter([( + "user_key".into(), + serde_json::json!("value"), + )])), + metadata_patch: None, + status: None, + }, + ) + .await; + let forged = sessions + .update_session( + session_id, + user_id.clone(), + astra_services::SessionUpdateRequestData { + title: None, + metadata: None, + metadata_patch: Some(serde_json::Map::from_iter([( + proof_key.into(), + serde_json::Value::Null, + )])), + status: None, + }, + ) + .await; + (Some(replacement), Some(forged)) + } else { + (None, None) + }; + let replay = sessions + .create_provider_session(first_identity.clone(), request.clone(), LimitCheck::Allowed) + .await; + let (old_replay, corrupt_replay) = if let Ok(first) = first.as_ref() { + // Existing deterministic creation intents remain readable across an + // upgrade. Replaying them must not invent ownership of old history. + sqlx::query( + "UPDATE agent_sessions SET metadata = '{}' WHERE user_id = ? AND session_id = ?", + ) + .bind(&user_id) + .bind(&first.session.session_id) + .execute(pool.get()) + .await + .unwrap(); + let old = sessions + .create_provider_session(first_identity.clone(), request.clone(), LimitCheck::Allowed) + .await; + sqlx::query("UPDATE agent_sessions SET metadata = ? WHERE user_id = ? AND session_id = ?") + .bind(serde_json::json!({(proof_key): null}).to_string()) + .bind(&user_id) + .bind(&first.session.session_id) + .execute(pool.get()) + .await + .unwrap(); + let corrupt = sessions + .create_provider_session(first_identity, request.clone(), LimitCheck::Allowed) + .await; + (Some(old), Some(corrupt)) + } else { + (None, None) + }; + // A stale preflight check must not bypass the authoritative transaction. + let second = sessions + .create_provider_session( + ProviderSessionCreationIdentity::from_principal(&principal, "second").unwrap(), + request.clone(), + LimitCheck::Allowed, + ) + .await; + let usage = governor.get_usage(&user_id).await; + let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM agent_sessions WHERE user_id = ?") + .bind(&user_id) + .fetch_one(pool.get()) + .await + .unwrap(); + governor + .set_limits( + &user_id, + ResourceLimits { + max_sessions_per_day: 2, + ..Default::default() + }, + ) + .await; + // The 255-character title column rejects this insert after quota reservation. + let failed_insert = sessions + .create_session_with_resource_quota( + user_id.clone(), + SessionCreateRequestData { + title: Some("x".repeat(1_024)), + ..request.clone() + }, + &governor, + ) + .await; + let after_rollback = governor.get_usage(&user_id).await; + let (ordinary, provider) = tokio::join!( + sessions.create_session_with_resource_quota(user_id.clone(), request.clone(), &governor), + sessions.create_provider_session( + ProviderSessionCreationIdentity::from_principal(&principal, "concurrent").unwrap(), + request, + LimitCheck::Allowed, + ), + ); + let after_race = governor.get_usage(&user_id).await; + let count_after_race: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM agent_sessions WHERE user_id = ?") + .bind(&user_id) + .fetch_one(pool.get()) + .await + .unwrap(); + pool.get().close().await; + sqlx::query(sqlx::AssertSqlSafe(format!("DROP DATABASE `{database}`"))) + .execute(&mut admin) + .await + .unwrap(); + let replacement = replacement + .expect("provider session created") + .expect("metadata replacement"); + assert!(original_proof.is_some()); + assert_eq!(replacement.metadata.get(proof_key), original_proof.as_ref()); + assert_eq!( + forged + .expect("provider session created") + .expect_err("scope proof removal must be rejected") + .0, + StatusCode::BAD_REQUEST + ); + assert!( + internal.is_ok(), + "core-only internal Session creation remains supported" + ); + assert_eq!( + unavailable.unwrap_err().0, + StatusCode::INTERNAL_SERVER_ERROR + ); + assert_eq!( + unavailable_sessions, 0, + "quota storage failure must not create a Session" + ); + assert!(first.unwrap().created); + assert!(!replay.unwrap().created); + let old_replay = old_replay.unwrap().unwrap(); + assert!(!old_replay.created); + assert!( + !old_replay.session.metadata.contains_key(proof_key), + "replay must not backfill scope" + ); + assert_eq!( + corrupt_replay.unwrap().unwrap_err().0, + StatusCode::FORBIDDEN + ); + assert_eq!(second.unwrap_err().0, StatusCode::TOO_MANY_REQUESTS); + assert_eq!( + usage.sessions_created, 1, + "replay must not consume quota twice" + ); + assert_eq!(count, 1, "quota rejection must not leave a Session"); + assert_eq!( + failed_insert.unwrap_err().0, + StatusCode::INTERNAL_SERVER_ERROR + ); + assert_eq!( + after_rollback.sessions_created, 1, + "failed creation releases its count" + ); + let race = [ordinary.map(|_| ()), provider.map(|_| ())]; + assert_eq!(race.iter().filter(|result| result.is_ok()).count(), 1); + for error in race.into_iter().filter_map(Result::err) { + assert_eq!(error.0, StatusCode::TOO_MANY_REQUESTS); + } + assert_eq!(after_race.sessions_created, 2); + assert_eq!( + count_after_race, 2, + "ordinary and provider creation share the quota lock" + ); +} diff --git a/docs/design/session-observability.md b/docs/design/session-observability.md index 9fcba273da..9f35a6fa48 100644 --- a/docs/design/session-observability.md +++ b/docs/design/session-observability.md @@ -272,6 +272,19 @@ Physical page counts, page high watermarks, and page lag are not exposed. The committed projection may honestly lag canonical context until its evidence has been materialized and the contiguous cursor promoted. +A Session fork installs a validated shared prefix and child canonical head without +copying parent display rows or allocating historical child runs. The child's +first own run advances from turn N to N+1 through the existing canonical terminal +transaction, which commits Context, raw transcript and terminal status together. +Repeated activation cannot rewind a continued child. + +Nested forks use the existing canonical lineage owner. Each ancestor read remains +within its validated owner, isolation domain, branch, generation and pinned +sequence. Deleting an intermediate Session removes its live head and execution +authority while existing fork records and pins retain history needed by prepared +or active descendants. Bounded bottom-up cleanup releases those records and +node-local payloads after the last dependent descendant is removed. + ## One-shot local content privacy `astra --no-journal-content chat ...` (or `ASTRA_JOURNAL_CONTENT_REDACT=1`) diff --git a/docs/guides/testing.md b/docs/guides/testing.md index 6b968ae4b8..3cd99f4e85 100644 --- a/docs/guides/testing.md +++ b/docs/guides/testing.md @@ -207,6 +207,11 @@ requires the item never to start. Use `require_board_before_execution` to reject unbound pre-reads and require successful evidence under the delivered assignment. Parallel batch position alone does not prove cancellation/execution ordering. +The offline `multi_agent_binding_http_e2e` journeys exercise sessionless runtime +execution, real model/skill HTTP, and canonical SSE projection. They do not +exercise inbound HTTP authentication or persistent Session creation/admission; +those boundaries require the real database consumer journey below. + ## MOI consumer compatibility with a fresh Astra database `Test: MOI fresh-Astra compatibility` starts real MatrixOne and the real Astra @@ -289,6 +294,48 @@ using Astra request types or fetching a moving MOI branch. It checks: - a valid turn still completes on that session after the rejected requests, without duplicate inference. +The additive `moi_fresh_astra/forks.rs` contract uses the same independent +signer and real router, with its own `astra_test_moi_fork_` database. +It exercises signed prepare, continuation activation, and exact committed +root-run cursor lookup, then continues the parent, two siblings, and a nested +fork through ordinary chat. Signature path and replay controls use the supported +cursor route; legacy list/read/writer-activation routes retain their +existing authentication rather than admitting request-signed provider tokens. +The finite model inventory checks the actual inherited input and excludes later +parent or sibling turns. Activation replay must preserve a continued child's +head; activation alone grants no writer lease and makes no inference call. +Negative controls cover signature/body/path tampering, nonce replay, workspace +scope isolation, idempotency conflicts, non-root runs, incomplete projection, +and unavailable historical roots. The earlier consumer baseline is unchanged. +The same fork journey deletes its middle Session, then its ancestor, and proves +the descendant can still materialize and continue from the exact pinned history +after bounded cleanup. Deleted Sessions retain no live head or execution state. +Deleting the final descendant releases fork metadata, pins and manifest rows, +including the node-local payloads stored by the current canonical owner. + +`moi_fresh_astra/fork_auth_compat.rs` checks that real runtime and Edge user +tokens retain the existing prepare and writer-activation contract at a +canonical head before Server projection. These legacy credentials cannot use +the new provider cursor or continuation-activation controls. A custom +request-only provider authenticator without mechanism evidence cannot bypass +scope or retry legacy authentication. The canonical verifier also proves the +legacy/HMAC mechanism distinction and single nonce consumption. + +MOI's new provider fork integration uses request-bound HMAC authentication. +Those principals own Sessions under `provider_authorized::`; +legacy MOI user tokens use the separate `external_authorized:moi:` +owner namespace and retain their existing subject-owned Session contract. +Legacy credentials cannot read or control a request-signed provider's fork. +This integration does not change legacy APIs into workspace-scoped APIs. + +New request-signed provider Sessions and forks enforce scope at common Run +admission and access. Existing ordinary Sessions without a scope binding retain +their previous chat behavior, but cannot serve as a source for a new provider +fork. A later Run never claims that Session's historical scope. A malformed +binding or an existing fork manifest cannot use the ordinary-session fallback. +The fresh-database fixtures do not establish an existing-database migration or +repair older ordinary Sessions' pre-existing scope behavior. + Buffered and live consumers share the pinned MOI SSE framing rules, including LF/CRLF, multiline data, EOF frames and stopping at `[DONE]`. Negative controls reject an early `[DONE]` even when a terminal follows it. MCP and Skill callbacks diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index 4d57caadb2..09dde9226c 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -277,6 +277,83 @@ Physical page references/hashes are no longer part of this response. Work transcript responses retain their committed cursor and per-item commitment fields; a newer canonical context head does not certify transcript publication. +### Session forks + +The operations below retain ordinary Astra owner authentication. Provider access +to the new run-cursor and continuation-activation operations requires +request-bound HMAC authentication, binding the exact method, path, route +template, body digest, and fresh request ID. Scope is verified against the +source Session's server-authenticated binding or its retained fork manifest, +not a display message or the latest Run. The child retains that scope. + +Existing runtime/Edge user tokens retain the original prepare and writer- +activation contract, including canonical heads before Server transcript +projection. They cannot use the new provider cursor or continuation controls, +or access request-signed Sessions in the distinct provider owner namespace. +Provider authenticators without explicit mechanism provenance are rejected. +Invalid request-HMAC credentials never fall back to legacy authentication. + +New request-signed Sessions and forks enforce scope on subsequent Run admission +and access. Older ordinary Sessions without a binding retain their previous +chat behavior but cannot become new provider fork sources. A later Run does not +claim their historical scope; malformed bindings and fork manifests cannot use +the ordinary-session compatibility path. + +`GET /sessions/{session_id}/runs/{run_id}/cursor` returns a `SessionCursorV1` +for that completed root conversation run. It verifies the run's committed +transcript markers against its retained canonical manifest. It never substitutes +the current Session head, counts display messages, or reconstructs a legacy +cursor from text. Canonical terminal settlement commits Context and raw transcript +with terminal status; a client querying an incomplete run may receive +`session_fork_cursor_not_ready` and retry this read. + +`POST /sessions/{session_id}/forks` accepts +`{idempotency_key, expected_parent_cursor, reason}` and returns the prepared +`SessionForkManifestV1`. The key identifies the immutable request; changing the +request under the same key conflicts. + +`POST /sessions/{session_id}/forks/{fork_id}/activate-for-continuation` accepts +`{}` and returns `{manifest, child_head}`. It activates the shared conversation +prefix without acquiring or returning a writer lease. The next ordinary +`/chat/stream` request uses `manifest.child_key.session_id` and admits fresh +Run-owned Context authority. Retrying activation returns the current child head and never +rewinds a continued child. This endpoint does not replay tools or inherit runs, +writers, approvals, mailboxes, invocations, filesystem state, or model credentials. +The existing `GET /sessions/{session_id}/forks` list, +`GET /sessions/{session_id}/forks/{fork_id}` read, +and `POST /sessions/{session_id}/forks/{fork_id}/activate` retain their existing +authentication contract; request-signed provider credentials do not grant +access to them. `/activate` retains its writer-lease response. + +`POST /sessions/{session_id}/forks/{fork_id}/abort` retains +`{detail, grace_seconds}` and additionally accepts request-signed provider +credentials for the matching prepared fork. It cancels preparation without +granting execution authority; active forks cannot be aborted. Interrupted +provider preparations expire after 24 hours and maintenance reclaims the +unused child and pin. Replaying a confirmed aborted operation returns +`session_fork_aborted`, including after child cleanup. Unknown outcomes must +still be retried with the same idempotency key. + +Fork errors use the existing public error envelope and stable `error_code`: + +| HTTP | Code | Meaning | +| --- | --- | --- | +| 400 | `session_fork_invalid` | Invalid request | +| 413 | `session_fork_resource_limit` | Canonical context resource limit exceeded | +| 403 | `session_fork_forbidden` | Source is outside the authorized provider scope | +| 404 | `session_fork_not_found` | Scoped run or fork was not found | +| 409 | `session_fork_cursor_not_ready` | Run or committed transcript projection is incomplete | +| 409 | `session_fork_cursor_unavailable` | Exact retained canonical root is unavailable | +| 422 | `session_fork_cursor_unsupported` | Run is not an eligible root conversation commit | +| 409 | `session_fork_idempotency_mismatch` | Same key with a different request | +| 409 | `session_fork_writer_conflict` | Existing activation's writer belongs to another actor | +| 409 | `session_fork_conflict` | Fork lifecycle or context conflict | +| 410 | `session_fork_aborted` | Preparation was cancelled and can no longer activate | +| 429 | `per_user_daily_session_quota` | Daily Session creation quota exceeded | +| 429 | `per_user_concurrent_session_quota` | Concurrent Session quota exceeded | +| 429 | `per_user_daily_token_quota` | Daily token quota exceeded | +| 503 | `session_fork_dependency_unavailable` | Dependency or stored-state failure; internal details are not returned | + ### PUT /sessions/{session_id} ### POST /sessions/{session_id}/close diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 8ce116ac2a..c0e20dbb6f 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -102,11 +102,17 @@ trust store to the Server. Restart the process after changing its trust settings - `ASTRA_JWT_ALGORITHM` (default `HS256`) - `ASTRA_JWT_ACCESS_TTL_MINUTES` (default `10080` in code; production should override) - `ASTRA_JWT_REFRESH_TTL_DAYS` (code default `30`; `.env.example` selects `7`) -- `ASTRA_TOKEN_ENCRYPTION_KEY` (shared high-entropy secret for stored provider and browser session credentials; changing it makes provider credentials undecryptable and requires browser sign-in again) +- `ASTRA_TOKEN_ENCRYPTION_KEY` (shared high-entropy secret for stored provider and browser session credentials and request-signed Session scope bindings; changing it makes provider credentials undecryptable, requires browser sign-in again, and invalidates existing Session scope bindings) - `ASTRA_RUNTIME_ROOT_SECRET` ### Provider Request Auth +Replicas sharing request-signed provider Sessions must use the same +`ASTRA_TOKEN_ENCRYPTION_KEY`. New Sessions authenticate their provider/workspace +binding with this key. Key rotation does not turn an unverifiable binding into +an old unbound Session; access is rejected rather than adopting the new request's +workspace. + Provider-originated service requests are authenticated under `auth.provider_request_auth` in `server.toml`. Astra validates these request tokens locally; it does not call a provider callback endpoint during request admission.