fix: harden memory surface generation

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<dyn LlmClient>,
+    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<dyn LlmClient> {
+        Box::new(self.clone())
+    }
+
+    async fn stream(
+        &self,
+        request: Request,
+    ) -> Result<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + 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<Vec<server_api::SubjektivSurfacePoint>, HookError> {
+    ) -> Result<Vec<server_api::SubjektivSurfacePoint>, 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<server_api::SubjektivSurfaceMaterial>), 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<usize, HookError> {
+    // 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::SubjektivSurfaceMaterial>,
+    ) -> 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<Vec<String>>,
+    }
+
+    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<WorkspaceResponse, WorkspaceClientError> {
+            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<AtomicUsize>,
+    }
+
+    #[async_trait]
+    impl LlmClient for SurfaceEditorClient {
+        fn clone_boxed(&self) -> Box<dyn LlmClient> {
+            Box::new(self.clone())
+        }
+
+        async fn stream(
+            &self,
+            _request: Request,
+        ) -> Result<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + 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<dyn WorkspaceClient>,
+        calls: Arc<AtomicUsize>,
+    ) -> 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"));
+    }
+}

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<dyn Tool> = Arc::new(SubmitMemorySurfaceTool {
             state: state.clone(),
         });

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<Item = &'a SurfaceMaterial>,
+) -> Result<usize> {
+    let materials = materials.into_iter().collect::<Vec<_>>();
+    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<Option<SurfaceSnapshot>> {
+    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, &registration).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();

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.