From 504ea00ae3aa42a9894ff1169199642f89ecf583 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 2 Oct 2026 14:01:27 +0900 Subject: [PATCH] fix: harden memory surface generation --- .../builtin/memory_surface_lifecycle.rs | 590 +++++++++++++++++- .../feature/builtin/memory_surface_output.rs | 24 +- crates/workspace-server/src/subjektiv.rs | 289 ++++++++- docs/design/subjektiv-surface.md | 8 +- 4 files changed, 843 insertions(+), 68 deletions(-) diff --git a/crates/worker/src/feature/builtin/memory_surface_lifecycle.rs b/crates/worker/src/feature/builtin/memory_surface_lifecycle.rs index db7b3f9be..ca7c15d15 100644 --- a/crates/worker/src/feature/builtin/memory_surface_lifecycle.rs +++ b/crates/worker/src/feature/builtin/memory_surface_lifecycle.rs @@ -1,9 +1,12 @@ +use std::pin::Pin; use std::sync::Arc; use std::time::Duration; -use agen::llm_client::LlmClient; +use agen::llm_client::event::Event as LlmEvent; +use agen::llm_client::{ClientError, LlmClient, Request}; use arc_swap::ArcSwap; use async_trait::async_trait; +use futures_util::Stream; use crate::PromptCatalog; use crate::Scope; @@ -13,7 +16,7 @@ use crate::feature::background::{ }; use crate::feature::builtin::memory::WorkspaceMemoryBackendError; use crate::feature::builtin::memory_surface_output::{ - MemorySurfaceOutputFeature, MemorySurfaceOutputState, + MemorySurfaceOutputFeature, MemorySurfaceOutputState, submit_llm_tool_definition, }; use crate::feature::{ BackgroundTaskDeclaration, FeatureDescriptor, FeatureInstallContext, FeatureInstallError, @@ -31,6 +34,7 @@ const TASK_NAME: &str = "subjektiv-memory-surface"; const TASK_TIMEOUT: Duration = Duration::from_secs(300); const MAX_GENERATION_ATTEMPTS: usize = 2; const EDITOR_MAX_TURNS: u32 = 3; +const EDITOR_QUESTION: &str = "この主体が次の作業を始める際、毎回思い出しておくべきことは何か。継続的な制約、判断の前提、未解決事項、再発を避けたい教訓を、指定予算内でまとめる。"; #[derive(Clone)] pub(crate) struct SubjektivSurfaceLifecycleFeature { @@ -46,6 +50,45 @@ struct SubjektivSurfaceLifecycleTask { workspace_context: WorkerWorkspaceContext, } +const INPUT_BUDGET_ERROR_MARKER: &str = "surface_editor_input_budget_exceeded"; + +struct InputBudgetLlmClient { + inner: Box, + input_token_budget: usize, +} + +impl Clone for InputBudgetLlmClient { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone_boxed(), + input_token_budget: self.input_token_budget, + } + } +} + +#[async_trait] +impl LlmClient for InputBudgetLlmClient { + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + async fn stream( + &self, + request: Request, + ) -> Result> + Send>>, ClientError> + { + let estimated = estimated_editor_request_tokens(&request) + .map_err(|error| ClientError::Config(error.to_string()))?; + if estimated > self.input_token_budget { + return Err(ClientError::Config(format!( + "{INPUT_BUDGET_ERROR_MARKER}: request estimate {estimated} exceeds {}", + self.input_token_budget + ))); + } + self.inner.stream(request).await + } +} + impl SubjektivSurfaceLifecycleFeature { pub(crate) fn from_manifest( lifecycle_enabled: bool, @@ -107,6 +150,27 @@ impl FeatureModule for SubjektivSurfaceLifecycleFeature { } } +struct SurfaceEditorFailure { + error: HookError, + reason_code: &'static str, +} + +impl SurfaceEditorFailure { + fn input_budget(error: HookError) -> Self { + Self { + error, + reason_code: "input_budget_exhausted", + } + } + + fn editor(error: HookError) -> Self { + Self { + error, + reason_code: "editor_failed", + } + } +} + #[async_trait] impl FeatureBackgroundTask for SubjektivSurfaceLifecycleTask { async fn run( @@ -138,10 +202,10 @@ impl FeatureBackgroundTask for SubjektivSurfaceLifecycleTask { } else { match self.edit_surface(&generation, cancellation.clone()).await { Ok(points) => points, - Err(error) => { - self.record_failure(&generation.generation_id, "editor_failed") + Err(failure) => { + self.record_failure(&generation.generation_id, failure.reason_code) .await; - tracing::warn!(%error, "subject Memory surface editor failed"); + tracing::warn!(error = %failure.error, reason_code = failure.reason_code, "subject Memory surface editor failed"); return Ok(()); } } @@ -192,7 +256,7 @@ impl SubjektivSurfaceLifecycleTask { &self, generation: &server_api::SubjektivSurfacePrepareResponse, cancellation: BackgroundTaskCancellation, - ) -> Result, HookError> { + ) -> Result, SurfaceEditorFailure> { let language = self .manifest .feature @@ -200,27 +264,25 @@ impl SubjektivSurfaceLifecycleTask { .workspace_settings() .map(|settings| settings.language) .ok_or_else(|| { - HookError::new( + SurfaceEditorFailure::editor(HookError::new( HookErrorCategory::Internal, "surface editor requires bound Workspace Memory settings", - ) + )) })?; let system_prompt = self .prompts .load_full() .subjektiv_memory_surface_system(&language) - .map_err(|error| HookError::new(HookErrorCategory::Internal, error.to_string()))?; - let input = serde_json::to_string_pretty(&serde_json::json!({ - "question": "この主体が次の作業を始める際、毎回思い出しておくべきことは何か。継続的な制約、判断の前提、未解決事項、再発を避けたい教訓を、指定予算内でまとめる。", - "body_token_budget": generation.body_token_budget, - "input_token_budget": generation.input_token_budget, - "materials": generation.materials, - })) - .map_err(|error| HookError::new(HookErrorCategory::Internal, error.to_string()))?; - let output_state = MemorySurfaceOutputState::new( - generation.materials.clone(), - generation.body_token_budget, - ); + .map_err(|error| { + SurfaceEditorFailure::editor(HookError::new( + HookErrorCategory::Internal, + error.to_string(), + )) + })?; + let (input, editor_materials) = bounded_editor_input(system_prompt.as_str(), generation) + .map_err(SurfaceEditorFailure::input_budget)?; + let output_state = + MemorySurfaceOutputState::new(editor_materials, generation.body_token_budget); let features = FeatureRegistryBuilder::new() .with_module(MemorySurfaceOutputFeature::new(output_state.clone())); let cancellation_observer = cancellation.clone(); @@ -237,7 +299,10 @@ impl SubjektivSurfaceLifecycleTask { run_id: uuid::Uuid::now_v7(), }, manifest: self.manifest.clone(), - client: self.client.clone_boxed(), + client: Box::new(InputBudgetLlmClient { + inner: self.client.clone_boxed(), + input_token_budget: generation.input_token_budget, + }), system_prompt, input, cache_key: None, @@ -256,17 +321,23 @@ impl SubjektivSurfaceLifecycleTask { ) .await; if let Err(error) = result { - return Err(HookError::new( - HookErrorCategory::Internal, - format!("surface editor Internal Worker failed: {}", error.source), - )); + let message = format!("surface editor Internal Worker failed: {}", error.source); + let error = HookError::new(HookErrorCategory::Internal, message.clone()); + return Err(if message.contains(INPUT_BUDGET_ERROR_MARKER) { + SurfaceEditorFailure::input_budget(error) + } else { + SurfaceEditorFailure::editor(error) + }); } - output_state.submitted().ok_or_else(|| { - HookError::new( - HookErrorCategory::Internal, - "surface editor finished without valid structured output", - ) - }) + output_state + .submitted() + .ok_or_else(|| { + HookError::new( + HookErrorCategory::Internal, + "surface editor finished without valid structured output", + ) + }) + .map_err(SurfaceEditorFailure::editor) } async fn record_failure(&self, generation_id: &str, reason_code: &str) { @@ -283,6 +354,463 @@ impl SubjektivSurfaceLifecycleTask { } } +fn bounded_editor_input( + system_prompt: &str, + generation: &server_api::SubjektivSurfacePrepareResponse, +) -> Result<(String, Vec), HookError> { + let mut materials = generation.materials.clone(); + loop { + let input = serde_json::to_string(&serde_json::json!({ + "question": EDITOR_QUESTION, + "body_token_budget": generation.body_token_budget, + "input_token_budget": generation.input_token_budget, + "materials": &materials, + })) + .map_err(|error| HookError::new(HookErrorCategory::Internal, error.to_string()))?; + let request = Request::new() + .system(system_prompt) + .user(input.clone()) + .tool(submit_llm_tool_definition()); + let estimated_tokens = estimated_editor_request_tokens(&request)?; + if estimated_tokens <= generation.input_token_budget { + if materials.is_empty() && !generation.materials.is_empty() { + return Err(HookError::new( + HookErrorCategory::Internal, + format!( + "surface editor fixed prompt and material overhead left no grounded material within the {} token input budget", + generation.input_token_budget + ), + )); + } + return Ok((input, materials)); + } + if materials.pop().is_none() { + return Err(HookError::new( + HookErrorCategory::Internal, + format!( + "surface editor fixed prompt, JSON framing, and tool schema exceed the {} token input budget (estimated {estimated_tokens})", + generation.input_token_budget + ), + )); + } + } +} + +fn estimated_editor_request_tokens(request: &Request) -> Result { + // Serialize the actual prompt-bearing Agen request fields. This is the + // stable provider-independent boundary; provider-specific HTTP envelopes + // are intentionally outside the configured editor-input budget. + let request = serde_json::to_vec(&serde_json::json!({ + "system_prompt": &request.system_prompt, + "items": &request.items, + "tools": &request.tools, + })) + .map_err(|error| HookError::new(HookErrorCategory::Internal, error.to_string()))?; + Ok(request.len().saturating_add(3) / 4) +} + fn surface_hook_error(error: WorkspaceMemoryBackendError) -> HookError { HookError::new(HookErrorCategory::Internal, error.to_string()) } + +#[cfg(test)] +mod tests { + use std::pin::Pin; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex}; + + use agen::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent}; + use agen::llm_client::{ClientError, Request}; + use futures::Stream; + use memory::extract::CandidateKind; + + use super::*; + use crate::feature::FeatureId; + use crate::feature::background::FeatureBackgroundTaskRegistryBuilder; + use crate::hook::HookInvocationContext; + use crate::worker::{WorkspaceClientError, WorkspaceRequest, WorkspaceResponse}; + + fn material(id: &str, body_bytes: usize) -> server_api::SubjektivSurfaceMaterial { + server_api::SubjektivSurfaceMaterial { + memory_id: id.into(), + revision: 1, + kind: CandidateKind::Constraint, + body_md: "x".repeat(body_bytes), + why_useful: "budget fixture".into(), + staleness: None, + } + } + + fn generation( + input_token_budget: usize, + materials: Vec, + ) -> server_api::SubjektivSurfacePrepareResponse { + server_api::SubjektivSurfacePrepareResponse { + generation_id: "generation-1".into(), + store_revision: 1, + active_memory_count: materials.len(), + materials, + body_token_budget: 1_024, + input_token_budget, + per_kind_limit: 8, + total_material_limit: 24, + } + } + + fn editor_request(system_prompt: &str, input: &str) -> Request { + Request::new() + .system(system_prompt) + .user(input) + .tool(submit_llm_tool_definition()) + } + + #[derive(Debug)] + struct SurfaceWorkspaceClient { + conflict_count: usize, + input_token_budget: usize, + prepare_calls: AtomicUsize, + publish_calls: AtomicUsize, + failure_reasons: Mutex>, + } + + impl SurfaceWorkspaceClient { + fn new(conflict_count: usize) -> Self { + Self { + conflict_count, + input_token_budget: 12_000, + prepare_calls: AtomicUsize::new(0), + publish_calls: AtomicUsize::new(0), + failure_reasons: Mutex::new(Vec::new()), + } + } + + fn with_input_token_budget(input_token_budget: usize) -> Self { + Self { + input_token_budget, + ..Self::new(0) + } + } + + fn response(value: server_api::SubjektivMemoryBackendResponse) -> WorkspaceResponse { + WorkspaceResponse { + status: 200, + body: serde_json::to_string(&value).unwrap(), + } + } + } + + impl WorkspaceClient for SurfaceWorkspaceClient { + fn workspace_id(&self) -> Option<&str> { + Some("workspace-1") + } + + fn kind(&self) -> &str { + "surface-lifecycle-test" + } + + fn is_available(&self) -> bool { + true + } + + fn execute( + &self, + request: WorkspaceRequest, + ) -> Result { + let request: server_api::SubjektivMemoryBackendRequest = + serde_json::from_str(request.body.as_deref().unwrap_or_default()) + .map_err(|error| WorkspaceClientError::Request(error.to_string()))?; + match request.operation { + server_api::SubjektivMemoryBackendOperation::PrepareSurface(_) => { + let call = self.prepare_calls.fetch_add(1, Ordering::SeqCst) + 1; + let mut prepared = + generation(self.input_token_budget, vec![material("memory-1", 32)]); + prepared.generation_id = format!("generation-{call}"); + prepared.store_revision = call as u64; + Ok(Self::response( + server_api::SubjektivMemoryBackendResponse::SurfacePrepared(prepared), + )) + } + server_api::SubjektivMemoryBackendOperation::PublishSurface(input) => { + assert_eq!(input.points.len(), 1); + let call = self.publish_calls.fetch_add(1, Ordering::SeqCst) + 1; + if call <= self.conflict_count { + return Ok(WorkspaceResponse { + status: 409, + body: "store generation moved".into(), + }); + } + Ok(Self::response( + server_api::SubjektivMemoryBackendResponse::SurfacePublished( + server_api::SubjektivSurfacePublishResponse { + snapshot_id: format!("surface-{call}"), + built_from_store_revision: call as u64, + empty: false, + }, + ), + )) + } + server_api::SubjektivMemoryBackendOperation::FailSurface(input) => { + self.failure_reasons.lock().unwrap().push(input.reason_code); + Ok(Self::response( + server_api::SubjektivMemoryBackendResponse::SurfaceFailed( + server_api::SubjektivSurfaceFailureResponse { + store_revision: self.prepare_calls.load(Ordering::SeqCst) as u64, + status: "failed".into(), + }, + ), + )) + } + _ => Err(WorkspaceClientError::Unavailable( + "surface lifecycle test only accepts surface operations".into(), + )), + } + } + } + + #[derive(Clone)] + struct SurfaceEditorClient { + calls: Arc, + } + + #[async_trait] + impl LlmClient for SurfaceEditorClient { + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + async fn stream( + &self, + _request: Request, + ) -> Result> + Send>>, ClientError> + { + let call = self.calls.fetch_add(1, Ordering::SeqCst); + let events = if call % 2 == 0 { + vec![ + LlmEvent::tool_use_start(0, format!("submit-{call}"), "SubmitMemorySurface"), + LlmEvent::tool_input_delta( + 0, + serde_json::json!({ + "points": [{ + "body_md": "- Keep the confirmed constraint.", + "memory_refs": [{"memory_id": "memory-1", "revision": 1}] + }] + }) + .to_string(), + ), + LlmEvent::tool_use_stop(0), + LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] + } else { + vec![ + LlmEvent::text_block_start(0), + LlmEvent::text_delta(0, "submitted"), + LlmEvent::text_block_stop(0, None), + LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] + }; + Ok(Box::pin(futures::stream::iter(events.into_iter().map(Ok)))) + } + } + + fn test_manifest() -> WorkerManifest { + let mut manifest = WorkerManifest::from_toml( + r#" +[worker] +name = "surface-lifecycle-test" +scope = "main" + +[model] +scheme = "anthropic" +model_id = "test-model" + +[engine] + +[[scope.allow]] +target = "/surface-lifecycle-test" +permission = "write" +"#, + ) + .unwrap(); + manifest.feature.memory.profile.enabled = true; + manifest + .feature + .memory + .bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot { + workspace_id: "workspace-1".into(), + settings_revision: 1, + language: "English".into(), + }) + .unwrap(); + manifest + } + + fn test_task( + workspace_client: Arc, + calls: Arc, + ) -> SubjektivSurfaceLifecycleTask { + SubjektivSurfaceLifecycleTask { + workspace_client, + manifest: test_manifest(), + client: Box::new(SurfaceEditorClient { calls }), + prompts: Arc::new(ArcSwap::from(PromptCatalog::builtins_only().unwrap())), + workspace_context: WorkerWorkspaceContext::no_workspace(), + } + } + + async fn run_task(task: SubjektivSurfaceLifecycleTask) { + let declaration = BackgroundTaskDeclaration::worker_managed( + TASK_NAME, + "surface lifecycle integration test", + ); + let mut spec = BackgroundTaskSpec::single_flight(declaration, Duration::from_secs(5)); + spec.trigger = BackgroundTaskTrigger::RunCommitted; + let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); + builder + .register(FeatureId::builtin("surface-lifecycle-test"), spec, task) + .unwrap(); + let registry = builder.build(); + registry + .start_run_committed(HookInvocationContext { + worker_id: "worker-1".into(), + session_id: "session-1".into(), + session_revision: 1, + run_id: Some("run-1".into()), + ..Default::default() + }) + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if !registry.diagnostics().is_empty() { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("surface lifecycle task should finish"); + registry.shutdown().await.unwrap(); + } + + #[test] + fn editor_input_budget_counts_prompt_json_question_and_tool_schema() { + let generation = generation( + 4_000, + vec![material("memory-1", 9_000), material("memory-2", 9_000)], + ); + let (input, selected) = bounded_editor_input("system policy", &generation).unwrap(); + + assert_eq!(selected.len(), 1, "deterministically trims the tail"); + assert_eq!(selected[0].memory_id, "memory-1"); + assert!(input.contains(EDITOR_QUESTION)); + assert!( + estimated_editor_request_tokens(&editor_request("system policy", &input)).unwrap() + <= generation.input_token_budget + ); + let materials_only_tokens = (serde_json::to_vec(&selected) + .unwrap() + .len() + .saturating_add(3)) + / 4; + assert!( + estimated_editor_request_tokens(&editor_request("system policy", &input)).unwrap() + > materials_only_tokens, + "full accounting must include framing and the SubmitMemorySurface schema" + ); + } + + #[tokio::test] + async fn actual_editor_request_is_rejected_before_provider_when_over_budget() { + let calls = Arc::new(AtomicUsize::new(0)); + let client = InputBudgetLlmClient { + inner: Box::new(SurfaceEditorClient { + calls: calls.clone(), + }), + input_token_budget: 32, + }; + let request = editor_request("system policy", &"x".repeat(1_000)); + + let error = match client.stream(request).await { + Err(error) => error, + Ok(_) => panic!("over-budget request reached the provider"), + }; + assert!(error.to_string().contains(INPUT_BUDGET_ERROR_MARKER)); + assert_eq!(calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn correction_turns_remain_subject_to_editor_input_budget() { + let calls = Arc::new(AtomicUsize::new(0)); + let client = InputBudgetLlmClient { + inner: Box::new(SurfaceEditorClient { + calls: calls.clone(), + }), + input_token_budget: 2_000, + }; + let _stream = client + .stream(editor_request("system policy", "small initial request")) + .await + .unwrap(); + let request = + editor_request("system policy", "small initial request").assistant("x".repeat(12_000)); + let error = match client.stream(request).await { + Err(error) => error, + Ok(_) => panic!("over-budget request reached the provider"), + }; + + assert!(error.to_string().contains(INPUT_BUDGET_ERROR_MARKER)); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn lifecycle_rebuilds_once_after_a_store_generation_conflict() { + let workspace = Arc::new(SurfaceWorkspaceClient::new(1)); + let editor_calls = Arc::new(AtomicUsize::new(0)); + run_task(test_task(workspace.clone(), editor_calls.clone())).await; + + assert_eq!(workspace.prepare_calls.load(Ordering::SeqCst), 2); + assert_eq!(workspace.publish_calls.load(Ordering::SeqCst), 2); + assert!(workspace.failure_reasons.lock().unwrap().is_empty()); + assert_eq!(editor_calls.load(Ordering::SeqCst), 4); + } + + #[tokio::test] + async fn lifecycle_bounds_conflict_retries_and_records_terminal_failure() { + let workspace = Arc::new(SurfaceWorkspaceClient::new(MAX_GENERATION_ATTEMPTS)); + let editor_calls = Arc::new(AtomicUsize::new(0)); + run_task(test_task(workspace.clone(), editor_calls.clone())).await; + + assert_eq!(workspace.prepare_calls.load(Ordering::SeqCst), 2); + assert_eq!(workspace.publish_calls.load(Ordering::SeqCst), 2); + assert_eq!(editor_calls.load(Ordering::SeqCst), 4); + assert_eq!( + workspace.failure_reasons.lock().unwrap().as_slice(), + ["publish_failed"] + ); + } + + #[tokio::test] + async fn lifecycle_records_input_budget_exhaustion_without_provider_or_publish() { + let workspace = Arc::new(SurfaceWorkspaceClient::with_input_token_budget(1)); + let editor_calls = Arc::new(AtomicUsize::new(0)); + run_task(test_task(workspace.clone(), editor_calls.clone())).await; + + assert_eq!(workspace.prepare_calls.load(Ordering::SeqCst), 1); + assert_eq!(workspace.publish_calls.load(Ordering::SeqCst), 0); + assert_eq!(editor_calls.load(Ordering::SeqCst), 0); + assert_eq!( + workspace.failure_reasons.lock().unwrap().as_slice(), + ["input_budget_exhausted"] + ); + } + + #[test] + fn editor_input_budget_never_relabels_trimmed_nonempty_materials_as_empty() { + let generation = generation(1, vec![material("memory-1", 32)]); + let error = bounded_editor_input("system policy", &generation).unwrap_err(); + assert!(error.to_string().contains("input budget")); + } +} diff --git a/crates/worker/src/feature/builtin/memory_surface_output.rs b/crates/worker/src/feature/builtin/memory_surface_output.rs index 4b3233364..2527d0cd5 100644 --- a/crates/worker/src/feature/builtin/memory_surface_output.rs +++ b/crates/worker/src/feature/builtin/memory_surface_output.rs @@ -140,13 +140,27 @@ impl FeatureModule for MemorySurfaceOutputFeature { } } +fn submit_input_schema() -> serde_json::Value { + serde_json::to_value(schemars::schema_for!(SubmitMemorySurfaceParams)) + .unwrap_or_else(|_| serde_json::json!({})) +} + +pub(crate) fn submit_llm_tool_definition() -> agen::llm_client::ToolDefinition { + agen::llm_client::ToolDefinition::new(SUBMIT_TOOL) + .description(SUBMIT_DESCRIPTION) + .input_schema(submit_input_schema()) +} + fn submit_definition(state: MemorySurfaceOutputState) -> ToolDefinition { Arc::new(move || { - let schema = serde_json::to_value(schemars::schema_for!(SubmitMemorySurfaceParams)) - .unwrap_or_else(|_| serde_json::json!({})); - let meta = ToolMeta::new(SUBMIT_TOOL) - .description(SUBMIT_DESCRIPTION) - .input_schema(schema); + let definition = submit_llm_tool_definition(); + let meta = ToolMeta::new(definition.name) + .description( + definition + .description + .expect("SubmitMemorySurface description is fixed"), + ) + .input_schema(definition.input_schema); let tool: Arc = Arc::new(SubmitMemorySurfaceTool { state: state.clone(), }); diff --git a/crates/workspace-server/src/subjektiv.rs b/crates/workspace-server/src/subjektiv.rs index 3b5a8708a..749ea31b1 100644 --- a/crates/workspace-server/src/subjektiv.rs +++ b/crates/workspace-server/src/subjektiv.rs @@ -917,19 +917,9 @@ INSERT INTO surface_generation_state ( ) SELECT subjects.subject_id, subjects.store_revision, - CASE WHEN EXISTS ( - SELECT 1 FROM surface_snapshots current - WHERE current.subject_id = subjects.subject_id - AND current.built_from_store_revision = subjects.store_revision - ) THEN 'ready' ELSE 'dirty' END, - ( - SELECT current.snapshot_id FROM surface_snapshots current - WHERE current.subject_id = subjects.subject_id - AND current.built_from_store_revision = subjects.store_revision - ORDER BY current.created_at ASC, current.snapshot_id ASC - LIMIT 1 - ), + 'dirty', NULL, + 'legacy_snapshot_requires_regeneration', subjects.updated_at FROM subjects WHERE EXISTS ( @@ -1729,7 +1719,7 @@ impl SubjektivStore { } let (body_md, memory_refs) = validate_surface_points(&materials, &points)?; if let Some(existing) = - surface_snapshot_for_revision(transaction, subject_id, store_revision)? + ready_surface_snapshot_for_revision(transaction, subject_id, store_revision)? { if existing.body_md == body_md && existing.memory_refs == memory_refs { return Ok(existing); @@ -1959,29 +1949,85 @@ fn select_surface_materials( } } - let mut selected = Vec::new(); - let mut payload_bytes = 2usize; // JSON array delimiters. + let populated_kind_count = buckets.iter().filter(|bucket| !bucket.is_empty()).count(); + let per_kind_token_budget = SURFACE_MATERIAL_TOKEN_BUDGET / populated_kind_count.max(1); + let mut ordered = Vec::new(); for offset in 0..SURFACE_PER_KIND_LIMIT { - for bucket in &buckets { - let Some(material) = bucket.get(offset) else { - continue; - }; - if selected.len() == SURFACE_TOTAL_MATERIAL_LIMIT { - return Ok((active_memory_count, selected)); + for (bucket_index, bucket) in buckets.iter().enumerate() { + if let Some(material) = bucket.get(offset) { + ordered.push((bucket_index, material.clone())); } - let material_bytes = serde_json::to_vec(material)?.len().saturating_add(1); - if estimated_tokens_for_bytes(payload_bytes.saturating_add(material_bytes)) - > SURFACE_MATERIAL_TOKEN_BUDGET - { - continue; - } - payload_bytes = payload_bytes.saturating_add(material_bytes); - selected.push(material.clone()); } } + + let mut accepted = vec![false; ordered.len()]; + let mut selected_count = 0usize; + let mut deferred = Vec::new(); + for (index, (bucket_index, material)) in ordered.iter().enumerate() { + if selected_count == SURFACE_TOTAL_MATERIAL_LIMIT { + break; + } + let kind_materials = ordered + .iter() + .zip(&accepted) + .filter(|((selected_bucket, _), accepted)| { + **accepted && selected_bucket == bucket_index + }) + .map(|((_, material), _)| material) + .chain(std::iter::once(material)); + let all_materials = ordered + .iter() + .zip(&accepted) + .filter(|(_, accepted)| **accepted) + .map(|((_, material), _)| material) + .chain(std::iter::once(material)); + if serialized_material_tokens(kind_materials)? <= per_kind_token_budget + && serialized_material_tokens(all_materials)? <= SURFACE_MATERIAL_TOKEN_BUDGET + { + accepted[index] = true; + selected_count += 1; + } else { + deferred.push(index); + } + } + + // Category shares are fairness reservations, not hard ceilings. Reclaim + // unused shares in the same canonical order after every populated category + // had its first opportunity. + for index in deferred { + if selected_count == SURFACE_TOTAL_MATERIAL_LIMIT { + break; + } + let material = &ordered[index].1; + let all_materials = ordered + .iter() + .zip(&accepted) + .filter(|(_, accepted)| **accepted) + .map(|((_, material), _)| material) + .chain(std::iter::once(material)); + if serialized_material_tokens(all_materials)? <= SURFACE_MATERIAL_TOKEN_BUDGET { + accepted[index] = true; + selected_count += 1; + } + } + + let selected = ordered + .into_iter() + .zip(accepted) + .filter_map(|((_, material), accepted)| accepted.then_some(material)) + .collect(); Ok((active_memory_count, selected)) } +fn serialized_material_tokens<'a>( + materials: impl IntoIterator, +) -> Result { + let materials = materials.into_iter().collect::>(); + Ok(estimated_tokens_for_bytes( + serde_json::to_vec(&materials)?.len(), + )) +} + fn surface_kind_index(kind: &CandidateKind) -> usize { match kind { CandidateKind::Preference => 0, @@ -2053,6 +2099,28 @@ fn validate_surface_points( Ok((body_md, refs)) } +fn ready_surface_snapshot_for_revision( + transaction: &Transaction<'_>, + subject_id: &str, + store_revision: u64, +) -> Result> { + let raw = transaction + .query_row( + "SELECT snapshots.snapshot_json + FROM surface_generation_state state + JOIN surface_snapshots snapshots + ON snapshots.subject_id = state.subject_id + AND snapshots.snapshot_id = state.snapshot_id + WHERE state.subject_id = ?1 + AND state.store_revision = ?2 + AND state.status = 'ready'", + params![subject_id, to_i64(store_revision)?], + |row| row.get::<_, String>(0), + ) + .optional()?; + raw.map(|raw| parse_surface_snapshot(&raw)).transpose() +} + fn surface_snapshot_for_revision( transaction: &Transaction<'_>, subject_id: &str, @@ -3237,6 +3305,19 @@ mod tests { "create subjektiv subject memory store", create_schema, )]; + static V3_MIGRATIONS: &[FeatureMigration] = &[ + FeatureMigration::new(1, "create subjektiv subject memory store", create_schema), + FeatureMigration::new( + 2, + "add immutable subject session attribution", + add_subject_session_attribution, + ), + FeatureMigration::new( + 3, + "add idempotent atomic candidate decision receipts", + add_candidate_decision_receipts, + ), + ]; fn role() -> SubjectRole { SubjectRole::new("workspace_companion").unwrap() @@ -3458,6 +3539,90 @@ mod tests { ); } + #[test] + fn legacy_surface_snapshot_migrates_as_stale_until_regenerated() { + let temp = tempfile::tempdir().unwrap(); + let root = temp.path().join("storage"); + let (subject_id, snapshot_id) = { + let manager = FeatureStorage::new(&root); + let workspace = manager.workspace("workspace-a").unwrap(); + let registration = workspace + .register(FeatureRegistration::new( + SUBJEKTIV_FEATURE_ID, + V3_MIGRATIONS, + )) + .unwrap(); + let store = SubjektivStore::open(&workspace, ®istration).unwrap(); + let subject = store.create_subject(role()).unwrap(); + let snapshot = SurfaceSnapshot { + schema_version: SUBJEKTIV_SCHEMA_VERSION, + id: "legacy-surface".into(), + subject_id: subject.id.clone(), + body_md: "Legacy summary without generation-policy evidence.".into(), + memory_refs: Vec::new(), + built_from_store_revision: subject.store_revision, + created_at: now(), + }; + store + .database + .try_transaction(|transaction| -> Result<()> { + transaction.execute( + "INSERT INTO surface_snapshots ( + subject_id, snapshot_id, built_from_store_revision, + snapshot_json, created_at + ) VALUES (?1, ?2, ?3, ?4, ?5)", + params![ + subject.id, + snapshot.id, + to_i64(snapshot.built_from_store_revision)?, + serde_json::to_string(&snapshot)?, + snapshot.created_at + ], + )?; + transaction.execute( + "INSERT INTO surface_snapshot_seals (subject_id, snapshot_id) + VALUES (?1, ?2)", + params![subject.id, snapshot.id], + )?; + Ok(()) + }) + .unwrap(); + (subject.id, snapshot.id) + }; + + let (_manager, _workspace, store) = open_store(&root, "workspace-a"); + let resident = store.resident_surface(&subject_id).unwrap(); + assert_eq!(resident.availability, SurfaceAvailability::Stale); + assert!(resident.snapshot.is_none()); + assert!( + store + .surface_snapshot(&subject_id, &snapshot_id) + .unwrap() + .is_some(), + "migration retains the legacy snapshot as history without injecting it" + ); + + let generation = store.prepare_surface_generation(&subject_id).unwrap(); + let regenerated = store + .publish_surface_generation(&subject_id, &generation.id, Vec::new()) + .unwrap(); + assert_ne!(regenerated.id, snapshot_id); + let ready = store.resident_surface(&subject_id).unwrap(); + assert_eq!(ready.availability, SurfaceAvailability::Ready); + assert_eq!(ready.snapshot.unwrap().id, regenerated.id); + assert!( + store + .surface_snapshot(&subject_id, &snapshot_id) + .unwrap() + .is_some(), + "regeneration appends instead of mutating legacy history" + ); + let retried = store + .publish_surface_generation(&subject_id, &generation.id, Vec::new()) + .unwrap(); + assert_eq!(retried.id, regenerated.id); + } + #[test] fn attributed_staging_is_atomic_and_retry_idempotent() { let temp = tempfile::tempdir().unwrap(); @@ -4812,6 +4977,72 @@ mod tests { assert!(estimated_tokens_for_bytes(serialized.len()) <= SURFACE_MATERIAL_TOKEN_BUDGET); } + #[test] + fn oversized_kind_cannot_starve_other_surface_categories() { + let temp = tempfile::tempdir().unwrap(); + let (_manager, _workspace, store) = open_store(temp.path(), "workspace-a"); + let subject = store.create_subject(role()).unwrap(); + + let mut oversized = draft("oversized preference", "category fairness fixture"); + oversized.kind = CandidateKind::Preference; + oversized.body_md = "x".repeat((SURFACE_MATERIAL_TOKEN_BUDGET - 100) * 4); + store.create_memory(&subject.id, oversized).unwrap(); + + for kind in [ + CandidateKind::Constraint, + CandidateKind::Decision, + CandidateKind::WorkingAssumption, + CandidateKind::OpenQuestion, + CandidateKind::Lesson, + ] { + let mut value = draft(&format!("{kind:?}"), "category fairness fixture"); + value.kind = kind; + store.create_memory(&subject.id, value).unwrap(); + } + + let generation = store.prepare_surface_generation(&subject.id).unwrap(); + for kind in [ + CandidateKind::Constraint, + CandidateKind::Decision, + CandidateKind::WorkingAssumption, + CandidateKind::OpenQuestion, + CandidateKind::Lesson, + ] { + assert!( + generation.materials.iter().any(|item| item.kind == kind), + "one large category must not exclude {kind:?}" + ); + } + assert!( + estimated_tokens_for_bytes(serde_json::to_vec(&generation.materials).unwrap().len()) + <= SURFACE_MATERIAL_TOKEN_BUDGET + ); + } + + #[test] + fn surface_material_selection_reclaims_unused_category_budget() { + let temp = tempfile::tempdir().unwrap(); + let (_manager, _workspace, store) = open_store(temp.path(), "workspace-a"); + let subject = store.create_subject(role()).unwrap(); + + let mut small = draft("small preference", "soft category budget fixture"); + small.kind = CandidateKind::Preference; + store.create_memory(&subject.id, small).unwrap(); + let mut large = draft("large constraint", "soft category budget fixture"); + large.kind = CandidateKind::Constraint; + large.body_md = "x".repeat(22_000); + store.create_memory(&subject.id, large).unwrap(); + + let generation = store.prepare_surface_generation(&subject.id).unwrap(); + assert_eq!(generation.materials.len(), 2); + assert_eq!(generation.materials[0].kind, CandidateKind::Preference); + assert_eq!(generation.materials[1].kind, CandidateKind::Constraint); + assert!( + estimated_tokens_for_bytes(serde_json::to_vec(&generation.materials).unwrap().len()) + <= SURFACE_MATERIAL_TOKEN_BUDGET + ); + } + #[test] fn surface_states_validation_races_and_idempotency_fail_closed() { let temp = tempfile::tempdir().unwrap(); diff --git a/docs/design/subjektiv-surface.md b/docs/design/subjektiv-surface.md index 41d46f311..457e2841c 100644 --- a/docs/design/subjektiv-surface.md +++ b/docs/design/subjektiv-surface.md @@ -12,7 +12,9 @@ Selection is deterministic: 2. each category is ordered by `updated_at DESC, memory_id ASC`; 3. at most 8 records are retained per category; 4. records are allocated round-robin by category, up to 24 records total; -5. the serialized material payload uses the repository's provider-independent token estimate, `ceil(UTF-8 bytes / 4)`, and is capped at 10,000 estimated tokens; the complete editor-input policy budget is 12,000 tokens. +5. the 10,000-token material budget is initially divided equally across categories as a fairness reservation; unused shares are then reclaimed in the same canonical round-robin order, so one large category cannot exclude another non-empty category or strand usable capacity; +6. both category and aggregate material bounds use the provider-independent estimate `ceil(UTF-8 bytes / 4)`; +7. before invoking the editor, the Worker measures the actual normalized Agen request fields containing the rendered system prompt, conversation items (including the fixed question and material JSON), and the typed `SubmitMemorySurface` tool definition/schema against the 12,000-token complete-input budget. It deterministically removes tail materials until the initial request fits, fails rather than publishing a false empty surface if none fit, and guards every subsequent correction request before it reaches the provider. Provider-specific HTTP envelope bytes are outside this normalized request budget. The clean-context editor uses the effective `builtin:subjektiv-memory-consolidation` profile model (`codex-oauth/gpt-5.6-luna`, medium reasoning). A generation permits three model turns (the initial attempt plus at most two correction turns). A store-generation conflict causes at most one complete re-read and regeneration. The published Markdown is capped at 1,024 tokens by the same estimate. @@ -20,7 +22,7 @@ Every editor point carries one or more exact `{memory_id, revision}` references. Publication is idempotent for identical output at one store revision. A different concurrent output cannot replace the first published surface. Confirmed Memory writes atomically mark the prior surface stale. Failed surface work never rolls back confirmed Memory or candidate disposition. -Resident injection loads only a `ready` snapshot built from the subject's exact current store revision. `un-generated`, `failed`, and `stale` states are omitted rather than represented as "no Memory". A successful empty generation is represented by a current ready snapshot with empty Markdown. Worker prompt materialization remains a start/restore boundary and does not rewrite recorded history. +Resident injection loads only a `ready` snapshot built from the subject's exact current store revision. `un-generated`, `failed`, and `stale` states are omitted rather than represented as "no Memory". A successful empty generation is represented by a current ready snapshot with empty Markdown. Snapshots created before this generation policy existed are retained as immutable history but migrate as `stale` and must be regenerated before injection. Worker prompt materialization remains a start/restore boundary and does not rewrite recorded history. ## Semantic fixture and verification @@ -43,4 +45,4 @@ A valid review fixture must: - cite only exact refs supplied with those materials; - distinguish mechanical citation validation from a human/model semantic review of the wording. -The repository tests exercise deterministic bounds, subject isolation, active-current-only selection, output/reference rejection, empty/un-generated/failed/stale states, generation races, parallel publication, correction/retraction staleness, and publication idempotency. The prompt fixture above is the manual semantic check for condition, negation, reservation, conflict, expiry, and duplicate editing behavior. +The repository tests exercise deterministic bounds and fair budget reclamation, subject isolation, active-current-only selection, complete normalized editor-request accounting, output/reference rejection, empty/un-generated/failed/stale states, legacy regeneration, generation races, bounded conflict retry, terminal failure recording, parallel publication, correction/retraction staleness, and publication idempotency. The prompt fixture above is the manual semantic check for condition, negation, reservation, conflict, expiry, and duplicate editing behavior.