fabro/lib/components/fabro-workflow/src/handler/agent.rs

use std::path::Path;
use std::sync::Arc;

use async_trait::async_trait;
use fabro_graphviz::graph::{Graph, Node};
use fabro_sandbox::RunSandbox;
use fabro_types::{StageModelUsage, StageTiming};
use pebble_agent::ToolMiddleware;
use pebble_coding_agent::extensions::HumanInputProvider;
pub(crate) use structured_output::extract_status_fields;
use tokio_util::sync::CancellationToken;

use super::llm::EffectiveRequestControls;
use super::structured_output::{
    self, OutputSchemaKind, StructuredOutputError, ValidatedStructuredOutput,
};
use super::{EngineServices, Handler, NodeTimeoutPolicy};
use crate::context::{Context, WorkflowContext, keys};
use crate::error::Error;
use crate::event::{Emitter, Event, StageScope};
use crate::interview_runtime::WorkflowHumanInput;
use crate::outcome::{BilledModelUsage, Outcome, OutcomeExt};

const LAST_FILE_ROUTING_EXTENSIONS: &[&str] = &["json", "md"];

/// Result from a `CodergenBackend` invocation.
#[allow(
    clippy::large_enum_variant,
    reason = "Text payload is the common case; Full(Box<Outcome>) is the rare alternative."
)]
pub enum CodergenResult {
    Text {
        text:              String,
        /// The stage's billing: for an agent, the whole session tree's
        /// tokens under the root's route.
        usage:             Option<BilledModelUsage>,
        /// `usage` split by model, when the backend billed subagents at
        /// their own models. Empty when `usage` is the one row.
        usage_by_model:    Vec<BilledModelUsage>,
        files_touched:     Vec<String>,
        last_file_touched: Option<String>,
        /// Active timing observed by the backend. The wall field is ignored by
        /// the executor on this hop; executor wall time remains authoritative.
        timing:            StageTiming,
    },
    Full(Box<Outcome>),
}

pub struct CodergenRunRequest<'a> {
    pub node:            &'a Node,
    pub prompt:          &'a str,
    pub context:         &'a Context,
    pub thread_id:       Option<&'a str>,
    pub emitter:         &'a Arc<Emitter>,
    pub sandbox:         &'a Arc<RunSandbox>,
    /// Tool hooks the stage's agent (and its subagents) run under.
    pub tool_middleware: Option<Arc<dyn ToolMiddleware>>,
    pub cancel_token:    CancellationToken,
    /// Where the agent's `ask_user` questions go.
    pub human_input:     Option<Arc<dyn HumanInputProvider>>,
}

pub struct OneShotRequest<'a> {
    pub node:          &'a Node,
    pub prompt:        &'a str,
    pub system_prompt: Option<&'a str>,
    pub emitter:       &'a Arc<Emitter>,
    pub stage_scope:   &'a StageScope,
    pub sandbox:       &'a Arc<RunSandbox>,
    pub cancel_token:  CancellationToken,
}

/// Emit the canonical `Event::Prompt` for a stage prompt and return the
/// resolved [`StageScope`] so the caller can keep building events scoped to
/// the same stage.
///
/// Both `AgentHandler` and `PromptHandler` build the same payload, so the
/// per-emit fallback rules — node-provided
/// `provider`/`model` overrides over run-level defaults, and the backend's
/// `EffectiveRequestControls` (or `Default::default()` when no backend is
/// attached) — live in one place.
pub(crate) fn emit_stage_prompt(
    services: &EngineServices,
    context: &Context,
    node: &Node,
    prompt: &str,
    mode: &str,
    backend: Option<&dyn CodergenBackend>,
) -> Result<StageScope, Error> {
    let prompt_provider = node
        .provider()
        .map(String::from)
        .or_else(|| Some(services.run.provider_id.to_string()));
    let prompt_model = node
        .model()
        .map(String::from)
        .or_else(|| Some(services.run.model.clone()));
    let stage_scope = StageScope::for_handler(context, &node.id);
    let request_controls = backend
        .map(|b| b.effective_request_controls(node))
        .transpose()?
        .unwrap_or_default();
    services.run.emitter.emit_scoped(
        &Event::Prompt {
            stage:            node.id.clone(),
            visit:            stage_scope.visit,
            text:             prompt.to_string(),
            mode:             Some(mode.to_string()),
            provider:         prompt_provider,
            model:            prompt_model,
            reasoning_effort: request_controls.reasoning_effort,
            speed:            request_controls.speed,
        },
        &stage_scope,
    );
    Ok(stage_scope)
}

/// Backend interface for LLM execution in codergen nodes.
#[async_trait]
pub trait CodergenBackend: Send + Sync {
    /// Run a multi-turn agent loop (the default codergen mode).
    async fn run(&self, request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error>;

    /// Run a single LLM call with no tools (one_shot mode).
    async fn one_shot(&self, _request: OneShotRequest<'_>) -> Result<CodergenResult, Error> {
        Err(Error::Validation(
            "one_shot mode not supported by this backend".into(),
        ))
    }

    async fn shutdown(&self, _emitter: &Arc<Emitter>) {}

    fn effective_request_controls(&self, _node: &Node) -> Result<EffectiveRequestControls, Error> {
        Ok(EffectiveRequestControls::default())
    }

    fn node_timeout_policy(&self, _node: &Node) -> NodeTimeoutPolicy {
        NodeTimeoutPolicy::ExecutorEnforced
    }
}

/// The default handler for LLM task nodes.
pub struct AgentHandler {
    backend: Option<Box<dyn CodergenBackend>>,
}

impl AgentHandler {
    #[must_use]
    pub fn new(backend: Option<Box<dyn CodergenBackend>>) -> Self {
        Self { backend }
    }
}

pub(crate) async fn validate_agent_output_sources(
    schema: &OutputSchemaKind,
    response_text: &str,
    sandbox: &Arc<RunSandbox>,
    last_file_touched: Option<&str>,
) -> Result<ValidatedStructuredOutput, StructuredOutputError> {
    if !matches!(schema, OutputSchemaKind::Routing) {
        return structured_output::validate_response_text(schema, response_text);
    }

    let initial_error = match structured_output::validate_response_text(schema, response_text) {
        Ok(validated) => return Ok(validated),
        Err(error) if error.allows_routing_fallback() => error,
        Err(error) => return Err(error),
    };

    let mut fallback_error = initial_error;
    if let Some(status_json) = read_sandbox_file(sandbox, "status.json").await {
        match structured_output::validate_response_text(schema, &status_json) {
            Ok(validated) => return Ok(validated),
            Err(error) if error.allows_routing_fallback() => {
                fallback_error = error;
            }
            Err(error) => return Err(error),
        }
    }

    if let Some(path) = last_file_touched {
        if let Some(routing_json) = read_last_file_routing_json(sandbox, path).await {
            return structured_output::validate_response_text(schema, &routing_json);
        }
    }

    Err(fallback_error)
}

async fn read_sandbox_file(sandbox: &Arc<RunSandbox>, path: &str) -> Option<String> {
    sandbox.read_file_text(path).await.ok()
}

/// Extract the terminal JSON object from the last-touched file when it has an
/// eligible extension. Does not check that the object contains routing fields;
/// callers validate that.
async fn read_last_file_routing_json(sandbox: &Arc<RunSandbox>, path: &str) -> Option<String> {
    let extension = Path::new(path).extension()?.to_str()?;
    if !LAST_FILE_ROUTING_EXTENSIONS
        .iter()
        .any(|allowed| extension.eq_ignore_ascii_case(allowed))
    {
        return None;
    }

    let contents = read_sandbox_file(sandbox, path).await?;
    structured_output::terminal_json_object(&contents).map(str::to_owned)
}

/// Truncate a string to at most `max_chars` characters (char-boundary safe).
pub(crate) fn truncate(s: &str, max_chars: usize) -> &str {
    if s.len() <= max_chars {
        s
    } else {
        &s[..s.floor_char_boundary(max_chars)]
    }
}

/// Shared simulate implementation for LLM-backed handlers (agent & prompt).
/// Produces a simulated outcome with standard context updates.
pub(crate) fn simulate_llm_handler(node: &Node) -> Outcome {
    let simulated_text = format!("[Simulated] Response for stage: {}", node.id);
    let mut outcome = Outcome::simulated(&node.id);
    outcome
        .context_updates
        .insert(keys::LAST_STAGE.to_string(), serde_json::json!(node.id));
    outcome.context_updates.insert(
        keys::LAST_RESPONSE.to_string(),
        serde_json::json!(truncate(&simulated_text, 200)),
    );
    outcome.context_updates.insert(
        keys::response_key(&node.id),
        serde_json::json!(&simulated_text),
    );
    outcome
}

#[async_trait]
impl Handler for AgentHandler {
    async fn shutdown(&self, emitter: &Arc<Emitter>) {
        if let Some(backend) = self.backend.as_ref() {
            backend.shutdown(emitter).await;
        }
    }

    async fn simulate(
        &self,
        node: &Node,
        _context: &Context,
        _graph: &Graph,
        _run_dir: &Path,
        _services: &EngineServices,
    ) -> Result<Outcome, Error> {
        Ok(simulate_llm_handler(node))
    }

    async fn execute(
        &self,
        node: &Node,
        context: &Context,
        graph: &Graph,
        _run_dir: &Path,
        services: &EngineServices,
    ) -> Result<Outcome, Error> {
        // 1. Build prompt (prepend fidelity preamble if present)
        let raw_prompt = node.prompt_or_label();
        let preamble = context.preamble();
        let prompt = if preamble.is_empty() {
            raw_prompt.to_string()
        } else {
            format!("{preamble}\n\n{raw_prompt}")
        };
        let output_schema = structured_output::parse_node_output_schema(node)?;
        let prompt = match output_schema.as_ref() {
            Some(schema) => schema.agent_prompt(&prompt),
            None => prompt,
        };

        let stage_scope = emit_stage_prompt(
            services,
            context,
            node,
            &prompt,
            StageModelUsage::MODE_AGENT,
            self.backend.as_deref(),
        )?;
        let human_input: Arc<dyn HumanInputProvider> = Arc::new(WorkflowHumanInput::new(
            Arc::clone(&services.interviewer),
            Arc::clone(&services.run.emitter),
            stage_scope.clone(),
            node.id.clone(),
            Arc::clone(&services.run.interview_blocker),
        ));

        // 3. Call LLM backend (agent loop)
        let thread_id = context.thread_id();
        let run_id = context.parsed_run_id()?;
        let tool_middleware: Option<Arc<dyn ToolMiddleware>> =
            services.run.hook_runner.as_ref().map(|hr| {
                Arc::new(fabro_hooks::WorkflowToolHookCallback {
                    hook_runner: Arc::clone(hr),
                    sandbox: Arc::clone(&services.run.sandbox),
                    run_id,
                    workflow_name: graph.name.clone(),
                    hook_execution_context: services.run.locations.hook_execution_context(),
                    node_id: node.id.clone(),
                }) as Arc<dyn ToolMiddleware>
            });
        let (
            response_text,
            stage_usage,
            stage_usage_by_model,
            backend_files_touched,
            last_file_touched,
            timing,
        ) = if let Some(backend) = &self.backend {
            let result = backend
                .run(CodergenRunRequest {
                    node,
                    prompt: &prompt,
                    context,
                    thread_id: thread_id.as_deref(),
                    emitter: &services.run.emitter,
                    sandbox: &services.run.sandbox,
                    tool_middleware,
                    cancel_token: services.run.cancel_token(),
                    human_input: Some(human_input),
                })
                .await;
            match result {
                Ok(CodergenResult::Full(outcome)) => return Ok(*outcome),
                Ok(CodergenResult::Text {
                    text,
                    usage,
                    usage_by_model,
                    files_touched,
                    last_file_touched,
                    timing,
                }) => (
                    text,
                    usage,
                    usage_by_model,
                    files_touched,
                    last_file_touched,
                    timing,
                ),
                Err(Error::Cancelled) => return Err(Error::Cancelled),
                Err(e) if e.is_retryable() => {
                    return Err(e);
                }
                Err(e) => {
                    return Ok(e.to_fail_outcome());
                }
            }
        } else {
            (
                format!("[Simulated] Response for stage: {}", node.id),
                None,
                Vec::new(),
                Vec::new(),
                None,
                StageTiming::default(),
            )
        };

        let response_model = stage_usage
            .as_ref()
            .map(|usage| usage.model_id().to_string())
            .or_else(|| node.model().map(String::from))
            .unwrap_or_default();
        let response_provider = node
            .provider()
            .map(String::from)
            .or_else(|| Some(services.run.provider_id.to_string()))
            .unwrap_or_default();
        services.run.emitter.emit_scoped(
            &Event::PromptCompleted {
                node_id:  node.id.clone(),
                response: response_text.clone(),
                model:    response_model,
                provider: response_provider,
                billing:  stage_usage.clone(),
            },
            &stage_scope,
        );

        // Build and write status
        let mut outcome = Outcome::success();
        outcome.notes = Some(format!("Stage completed: {}", node.id));
        outcome
            .context_updates
            .insert(keys::LAST_STAGE.to_string(), serde_json::json!(node.id));
        outcome.context_updates.insert(
            keys::LAST_RESPONSE.to_string(),
            serde_json::json!(truncate(&response_text, 200)),
        );
        outcome.context_updates.insert(
            keys::response_key(&node.id),
            serde_json::json!(&response_text),
        );

        if let Some(schema) = output_schema.as_ref() {
            if let Ok(validated) = validate_agent_output_sources(
                schema,
                &response_text,
                &services.run.sandbox,
                last_file_touched.as_deref(),
            )
            .await
            {
                structured_output::apply_validated_output(node, schema, &validated, &mut outcome);
            } else {
                let mut failed =
                    structured_output::exhausted_failure_outcome(node.output_retries());
                failed.timing = Some(timing);
                failed.usage = stage_usage;
                failed.usage_by_model = stage_usage_by_model;
                failed.files_touched = backend_files_touched;
                return Ok(failed);
            }
        } else {
            // 7b. Parse routing directives from response text, falling back to
            //     status.json written by the agent into the sandbox CWD, then to
            //     a terminal JSON object in an eligible last-written file.
            let found_in_response = extract_status_fields(&response_text, &mut outcome);
            if !found_in_response {
                let mut found_in_status_json = false;
                if let Some(status_json) =
                    read_sandbox_file(&services.run.sandbox, "status.json").await
                {
                    found_in_status_json = extract_status_fields(&status_json, &mut outcome);
                }
                if !found_in_status_json {
                    if let Some(ref path) = last_file_touched {
                        if let Some(routing_json) =
                            read_last_file_routing_json(&services.run.sandbox, path).await
                        {
                            extract_status_fields(&routing_json, &mut outcome);
                        }
                    }
                }
            }
        }
        outcome.usage = stage_usage;
        outcome.usage_by_model = stage_usage_by_model;
        outcome.files_touched = backend_files_touched;
        outcome.timing = Some(timing);

        Ok(outcome)
    }

    fn node_timeout_policy(&self, node: &Node) -> NodeTimeoutPolicy {
        self.backend
            .as_ref()
            .map_or(NodeTimeoutPolicy::ExecutorEnforced, |backend| {
                backend.node_timeout_policy(node)
            })
    }
}

#[cfg(test)]
#[expect(
    clippy::disallowed_methods,
    reason = "tests persist per-iteration state fixtures"
)]
mod tests {
    use std::sync::Arc;
    use std::time::Duration;

    use fabro_graphviz::graph::AttrValue;
    use fabro_store::{Database, RunDatabase, StageId};
    use fabro_types::{fixtures, test_support};
    use lithos_llm::types::{ReasoningEffort, Speed};
    use object_store::memory::InMemory;
    use tempfile::TempDir;

    use super::*;

    fn make_services() -> EngineServices {
        EngineServices::test_default()
    }

    fn test_store() -> Arc<Database> {
        Arc::new(fabro_store::test_support::test_database(
            Arc::new(InMemory::new()),
            "",
            Duration::from_millis(1),
            None,
        ))
    }

    async fn make_services_with_run_store() -> (
        EngineServices,
        RunDatabase,
        crate::event::StoreProgressLogger,
    ) {
        let store = test_store();
        let run_store = store.create_run(&fixtures::RUN_1).await.unwrap();
        seed_created(&run_store).await;
        let mut services = EngineServices::test_default();
        services.run = services
            .run
            .with_emitter(Arc::new(crate::event::Emitter::new(fixtures::RUN_1)))
            .with_run_store(run_store.clone().into());
        let logger = crate::event::StoreProgressLogger::new(run_store.clone());
        logger.register(services.run.emitter.as_ref());
        (services, run_store, logger)
    }

    async fn seed_created(run_store: &RunDatabase) {
        crate::event::append_event(
            run_store,
            &fixtures::RUN_1,
            &crate::event::Event::RunCreated {
                run_id:              fixtures::RUN_1,
                title:               None,
                settings:            serde_json::to_value(fabro_types::WorkflowSettings::default())
                    .unwrap(),
                graph:               serde_json::to_value(fabro_types::Graph::new("test")).unwrap(),
                workflow_source:     None,
                labels:              std::collections::BTreeMap::default(),
                source_directory:    None,
                workflow_slug:       None,
                workflow_version_id: None,
                target:              None,
                automation:          None,
                provenance:          test_support::test_run_provenance(),
                spec_blob:           None,
                git:                 None,
                fork_source_ref:     None,
                retried_from:        None,
                parent_id:           None,
                web_url:             None,
            },
        )
        .await
        .unwrap();
    }

    fn test_context() -> Context {
        let context = Context::new();
        context.set(
            crate::context::keys::INTERNAL_RUN_ID,
            serde_json::json!(fixtures::RUN_1.to_string()),
        );
        context
    }

    struct LastFileBackend {
        path: String,
    }

    #[async_trait]
    impl CodergenBackend for LastFileBackend {
        async fn run(&self, _request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
            Ok(CodergenResult::Text {
                text:              "Done writing results.".to_string(),
                usage_by_model:    Vec::new(),
                usage:             None,
                files_touched:     vec![self.path.clone()],
                last_file_touched: Some(self.path.clone()),
                timing:            StageTiming::default(),
            })
        }
    }

    async fn sandbox_with_file(path: &str, contents: &str) -> (TempDir, Arc<RunSandbox>) {
        let sandbox_dir = TempDir::new().unwrap();
        std::fs::write(sandbox_dir.path().join(path), contents).unwrap();
        let sandbox: Arc<RunSandbox> = Arc::new(
            fabro_sandbox::local_sandbox(sandbox_dir.path().to_path_buf())
                .await
                .unwrap(),
        );
        (sandbox_dir, sandbox)
    }

    async fn execute_with_last_file(path: &str, contents: &str) -> Outcome {
        let (_sandbox_dir, sandbox) = sandbox_with_file(path, contents).await;

        let handler = AgentHandler::new(Some(Box::new(LastFileBackend {
            path: path.to_string(),
        })));
        let node = Node::new("step");
        let context = test_context();
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        let mut services = EngineServices::test_default();
        services.run = services.run.with_sandbox(sandbox);

        handler
            .execute(&node, &context, &graph, tmp.path(), &services)
            .await
            .unwrap()
    }

    async fn validate_routing_with_last_file(
        path: &str,
        contents: &str,
    ) -> Result<ValidatedStructuredOutput, StructuredOutputError> {
        let (_sandbox_dir, sandbox) = sandbox_with_file(path, contents).await;

        validate_agent_output_sources(
            &OutputSchemaKind::Routing,
            "Done writing results.",
            &sandbox,
            Some(path),
        )
        .await
    }

    #[tokio::test]
    async fn codergen_handler_simulate() {
        let handler = AgentHandler::new(None);
        let node = Node::new("plan");
        let context = test_context();
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        let outcome = handler
            .simulate(&node, &context, &graph, tmp.path(), &make_services())
            .await
            .unwrap();
        assert_eq!(outcome.status, crate::outcome::StageOutcome::Succeeded);
        assert_eq!(outcome.notes.as_deref(), Some("[Simulated] plan"));
        assert_eq!(
            outcome.context_updates.get(keys::LAST_STAGE),
            Some(&serde_json::json!("plan"))
        );
        assert!(outcome.context_updates.contains_key(keys::LAST_RESPONSE));
        assert_eq!(
            outcome.context_updates.get(&keys::response_key("plan")),
            Some(&serde_json::json!("[Simulated] Response for stage: plan"))
        );
    }

    #[tokio::test]
    async fn codergen_handler_uses_already_rendered_prompt() {
        let handler = AgentHandler::new(None);
        let mut node = Node::new("plan");
        node.attrs.insert(
            "prompt".to_string(),
            AttrValue::String("Achieve: Build a feature".to_string()),
        );
        let context = test_context();
        let mut graph = Graph::new("test");
        graph.attrs.insert(
            "goal".to_string(),
            AttrValue::String("Build a feature".to_string()),
        );
        let tmp = TempDir::new().unwrap();
        let (services, run_store, logger) = make_services_with_run_store().await;

        handler
            .execute(&node, &context, &graph, tmp.path(), &services)
            .await
            .unwrap();
        logger.flush().await.unwrap();

        let state = run_store.state().await.unwrap();
        let node_state = state.stage(&StageId::new("plan", 1)).unwrap();
        assert_eq!(
            node_state.prompt.as_deref(),
            Some("Achieve: Build a feature")
        );
    }

    #[tokio::test]
    async fn codergen_handler_falls_back_to_label() {
        let handler = AgentHandler::new(None);
        let mut node = Node::new("work");
        node.attrs.insert(
            "label".to_string(),
            AttrValue::String("Do work".to_string()),
        );
        let context = test_context();
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();
        let (services, run_store, logger) = make_services_with_run_store().await;

        handler
            .execute(&node, &context, &graph, tmp.path(), &services)
            .await
            .unwrap();
        logger.flush().await.unwrap();

        let state = run_store.state().await.unwrap();
        let node_state = state.stage(&StageId::new("work", 1)).unwrap();
        assert_eq!(node_state.prompt.as_deref(), Some("Do work"));
    }

    #[tokio::test]
    async fn codergen_handler_context_updates() {
        let handler = AgentHandler::new(None);
        let node = Node::new("step");
        let context = test_context();
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        let outcome = handler
            .execute(&node, &context, &graph, tmp.path(), &make_services())
            .await
            .unwrap();

        assert_eq!(
            outcome.context_updates.get(keys::LAST_STAGE),
            Some(&serde_json::json!("step"))
        );
        assert!(outcome.context_updates.contains_key(keys::LAST_RESPONSE));
        assert_eq!(
            outcome.context_updates.get(&keys::response_key("step")),
            Some(&serde_json::json!("[Simulated] Response for stage: step"))
        );
    }

    #[tokio::test]
    async fn codergen_handler_falls_back_to_status_json_in_sandbox() {
        // Simulation mode returns text with no JSON directives, so the
        // handler should fall back to reading status.json from the sandbox CWD.
        let sandbox_dir = TempDir::new().unwrap();
        std::fs::write(
            sandbox_dir.path().join("status.json"),
            r#"{"outcome": "failed", "failure_reason": "tests failed"}"#,
        )
        .unwrap();

        let handler = AgentHandler::new(None);
        let node = Node::new("step");
        let context = test_context();
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        let mut services = EngineServices::test_default();
        services.run = services.run.with_sandbox(std::sync::Arc::new(
            fabro_sandbox::local_sandbox(sandbox_dir.path().to_path_buf())
                .await
                .unwrap(),
        ));

        let outcome = handler
            .execute(&node, &context, &graph, tmp.path(), &services)
            .await
            .unwrap();

        assert_eq!(outcome.status, crate::outcome::StageOutcome::Failed {
            retry_requested: false,
        });
        assert_eq!(outcome.failure_reason(), Some("tests failed"));
    }

    #[tokio::test]
    async fn codergen_handler_prefers_response_text_over_status_json() {
        // Backend returns response text with routing directives — status.json
        // in the sandbox should be ignored.
        struct DirectiveBackend;

        #[async_trait]
        impl CodergenBackend for DirectiveBackend {
            async fn run(&self, _request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
                Ok(CodergenResult::Text {
                    text:
                        r#"Done. {"outcome": "succeeded", "preferred_next_label": "approve"}"#
                            .to_string(),
                    usage_by_model:    Vec::new(),
                    usage:             None,
                    files_touched:     Vec::new(),
                    last_file_touched: None,
                    timing:            StageTiming::default(),
                })
            }
        }

        let sandbox_dir = TempDir::new().unwrap();
        std::fs::write(
            sandbox_dir.path().join("status.json"),
            r#"{"outcome": "failed", "failure_reason": "should be ignored"}"#,
        )
        .unwrap();

        let handler = AgentHandler::new(Some(Box::new(DirectiveBackend)));
        let node = Node::new("step");
        let context = test_context();
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        let mut services = EngineServices::test_default();
        services.run = services.run.with_sandbox(std::sync::Arc::new(
            fabro_sandbox::local_sandbox(sandbox_dir.path().to_path_buf())
                .await
                .unwrap(),
        ));

        let outcome = handler
            .execute(&node, &context, &graph, tmp.path(), &services)
            .await
            .unwrap();

        assert_eq!(outcome.status, crate::outcome::StageOutcome::Succeeded);
        assert_eq!(outcome.preferred_label.as_deref(), Some("approve"));
        assert!(outcome.failure.is_none());
    }

    #[tokio::test]
    async fn codergen_handler_copies_backend_timing_to_outcome() {
        struct TimingBackend;

        #[async_trait]
        impl CodergenBackend for TimingBackend {
            async fn run(&self, _request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
                Ok(CodergenResult::Text {
                    text:              "done".to_string(),
                    usage_by_model:    Vec::new(),
                    usage:             None,
                    files_touched:     Vec::new(),
                    last_file_touched: None,
                    timing:            StageTiming::new(0, 200, 300),
                })
            }
        }

        let handler = AgentHandler::new(Some(Box::new(TimingBackend)));
        let node = Node::new("step");
        let context = test_context();
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        let outcome = handler
            .execute(&node, &context, &graph, tmp.path(), &make_services())
            .await
            .unwrap();

        assert_eq!(outcome.timing, Some(StageTiming::new(0, 200, 300)));
    }

    #[tokio::test]
    async fn codergen_handler_extracts_status_from_last_file_touched() {
        let outcome = execute_with_last_file(
            "results.md",
            r#"# Results
{"context_updates": {"verified": "true"}}
"#,
        )
        .await;

        assert_eq!(outcome.status, crate::outcome::StageOutcome::Succeeded);
        assert_eq!(
            outcome.context_updates.get("verified"),
            Some(&serde_json::json!("true")),
        );
    }

    #[tokio::test]
    async fn codergen_handler_ignores_nonterminal_status_in_last_markdown_file() {
        let outcome = execute_with_last_file(
            "results.md",
            r#"{"outcome":"failed","failure_reason":"tests failed"}

All checks passed.
"#,
        )
        .await;

        assert_eq!(outcome.status, crate::outcome::StageOutcome::Succeeded);
        assert!(outcome.failure.is_none());
    }

    #[tokio::test]
    async fn codergen_handler_ignores_terminal_status_in_disallowed_last_file() {
        let outcome = execute_with_last_file(
            "command.rs",
            r#"{"outcome":"failed","failure_reason":"tests failed"}"#,
        )
        .await;

        assert_eq!(outcome.status, crate::outcome::StageOutcome::Succeeded);
        assert!(outcome.failure.is_none());
    }

    #[tokio::test]
    async fn codergen_handler_output_schema_routing_uses_status_json_fallback_when_response_has_no_json()
     {
        let sandbox_dir = TempDir::new().unwrap();
        std::fs::write(
            sandbox_dir.path().join("status.json"),
            r#"{"preferred_next_label": "review"}"#,
        )
        .unwrap();

        let handler = AgentHandler::new(None);
        let mut node = Node::new("step");
        node.attrs.insert(
            "output_schema".to_string(),
            AttrValue::String("routing".to_string()),
        );
        let context = test_context();
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        let mut services = EngineServices::test_default();
        services.run = services.run.with_sandbox(std::sync::Arc::new(
            fabro_sandbox::local_sandbox(sandbox_dir.path().to_path_buf())
                .await
                .unwrap(),
        ));

        let outcome = handler
            .execute(&node, &context, &graph, tmp.path(), &services)
            .await
            .unwrap();

        assert_eq!(outcome.status, crate::outcome::StageOutcome::Succeeded);
        assert_eq!(outcome.preferred_label.as_deref(), Some("review"));
    }

    #[tokio::test]
    async fn validated_routing_accepts_terminal_status_in_json_file_case_insensitively() {
        let validated = validate_routing_with_last_file(
            "results.JSON",
            "# Results\n\n{\"preferred_next_label\":\"review\"}\n",
        )
        .await
        .unwrap();

        assert_eq!(
            validated.value,
            serde_json::json!({"preferred_next_label": "review"}),
        );
    }

    #[tokio::test]
    async fn validated_routing_ignores_nonterminal_status_in_last_markdown_file() {
        let error = validate_routing_with_last_file(
            "results.md",
            "{\"outcome\":\"failed\",\"failure_reason\":\"tests failed\"}\nAll checks passed.",
        )
        .await
        .unwrap_err();

        assert_eq!(
            error.kind(),
            structured_output::StructuredOutputErrorKind::NoJsonObject,
        );
    }

    #[tokio::test]
    async fn validated_routing_ignores_terminal_status_in_disallowed_last_file() {
        let error = validate_routing_with_last_file(
            "command.rs",
            r#"{"outcome":"failed","failure_reason":"tests failed"}"#,
        )
        .await
        .unwrap_err();

        assert_eq!(
            error.kind(),
            structured_output::StructuredOutputErrorKind::NoJsonObject,
        );
    }

    #[tokio::test]
    async fn codergen_handler_output_schema_routing_rejects_malformed_response_before_status_json_fallback()
     {
        struct BadRoutingBackend;

        #[async_trait]
        impl CodergenBackend for BadRoutingBackend {
            async fn run(&self, _request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
                Ok(CodergenResult::Text {
                    text:              r#"{"suggested_next_ids": [1]}"#.to_string(),
                    usage_by_model:    Vec::new(),
                    usage:             None,
                    files_touched:     Vec::new(),
                    last_file_touched: None,
                    timing:            StageTiming::default(),
                })
            }
        }

        let sandbox_dir = TempDir::new().unwrap();
        std::fs::write(
            sandbox_dir.path().join("status.json"),
            r#"{"preferred_next_label": "should_not_use"}"#,
        )
        .unwrap();

        let handler = AgentHandler::new(Some(Box::new(BadRoutingBackend)));
        let mut node = Node::new("step");
        node.attrs.insert(
            "output_schema".to_string(),
            AttrValue::String("routing".to_string()),
        );
        node.attrs
            .insert("output_retries".to_string(), AttrValue::Integer(0));
        let context = test_context();
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        let mut services = EngineServices::test_default();
        services.run = services.run.with_sandbox(std::sync::Arc::new(
            fabro_sandbox::local_sandbox(sandbox_dir.path().to_path_buf())
                .await
                .unwrap(),
        ));

        let outcome = handler
            .execute(&node, &context, &graph, tmp.path(), &services)
            .await
            .unwrap();

        assert_eq!(outcome.status, crate::outcome::StageOutcome::Failed {
            retry_requested: false,
        });
        assert_eq!(
            outcome.failure_reason(),
            Some("output schema validation failed after 0 repair attempt(s)")
        );
        assert!(outcome.preferred_label.is_none());
    }

    #[tokio::test]
    async fn codergen_handler_custom_output_schema_updates_output_context_key() {
        struct CustomOutputBackend;

        #[async_trait]
        impl CodergenBackend for CustomOutputBackend {
            async fn run(&self, _request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
                Ok(CodergenResult::Text {
                    text:              r#"{"passed": true}"#.to_string(),
                    usage_by_model:    Vec::new(),
                    usage:             None,
                    files_touched:     Vec::new(),
                    last_file_touched: None,
                    timing:            StageTiming::default(),
                })
            }
        }

        let handler = AgentHandler::new(Some(Box::new(CustomOutputBackend)));
        let mut node = Node::new("audit");
        node.attrs.insert(
            "output_schema".to_string(),
            AttrValue::String(
                r#"{"type":"object","required":["passed"],"properties":{"passed":{"type":"boolean"}}}"#
                    .to_string(),
            ),
        );
        let context = test_context();
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        let outcome = handler
            .execute(&node, &context, &graph, tmp.path(), &make_services())
            .await
            .unwrap();

        assert_eq!(
            outcome.context_updates.get("output.audit"),
            Some(&serde_json::json!({"passed": true})),
        );
    }

    #[tokio::test]
    async fn codergen_handler_appends_output_schema_contract_to_prompt() {
        use std::sync::{Arc, Mutex};

        struct PromptCapturingBackend {
            captured_prompt: Arc<Mutex<Option<String>>>,
        }

        #[async_trait]
        impl CodergenBackend for PromptCapturingBackend {
            async fn run(&self, request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
                *self.captured_prompt.lock().unwrap() = Some(request.prompt.to_string());
                Ok(CodergenResult::Text {
                    text:              r#"{"passed": true}"#.to_string(),
                    usage_by_model:    Vec::new(),
                    usage:             None,
                    files_touched:     Vec::new(),
                    last_file_touched: None,
                    timing:            StageTiming::default(),
                })
            }
        }

        let captured = Arc::new(Mutex::new(None));
        let handler = AgentHandler::new(Some(Box::new(PromptCapturingBackend {
            captured_prompt: captured.clone(),
        })));

        let mut node = Node::new("audit");
        node.attrs.insert(
            "prompt".to_string(),
            AttrValue::String("Audit the result".to_string()),
        );
        node.attrs.insert(
            "output_schema".to_string(),
            AttrValue::String(
                r#"{"type":"object","required":["passed"],"properties":{"passed":{"type":"boolean"}}}"#
                    .to_string(),
            ),
        );
        let context = test_context();
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        handler
            .execute(&node, &context, &graph, tmp.path(), &make_services())
            .await
            .unwrap();

        let prompt = captured.lock().unwrap().clone().unwrap();
        assert!(
            prompt.starts_with("Audit the result\n\n"),
            "task prompt should come first, got: {prompt}"
        );
        assert!(
            prompt.contains("Fabro final-output contract"),
            "contract heading missing, got: {prompt}"
        );
        assert!(
            prompt.contains(
                "It applies only to your final response, not to intermediate tool calls."
            ),
            "contract should scope itself to the final response, got: {prompt}"
        );
        assert!(
            prompt.contains(r#""required":["passed"]"#),
            "contract should embed the resolved schema, got: {prompt}"
        );
        assert!(
            prompt.ends_with("Do not ask the user to provide or choose the output shape."),
            "contract should close the prompt, got: {prompt}"
        );
    }

    #[tokio::test]
    async fn codergen_handler_projects_provider_used_from_agent_session_events() {
        struct ProviderEventBackend;

        #[async_trait]
        impl CodergenBackend for ProviderEventBackend {
            async fn run(&self, request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
                let scope = StageScope::for_handler(request.context, &request.node.id);
                request.emitter.emit_scoped(
                    &crate::event::Event::AgentSessionActivated {
                        node_id:          request.node.id.clone(),
                        visit:            scope.visit,
                        session_id:       "session_123".to_string(),
                        thread_id:        None,
                        provider:         Some("openai".to_string()),
                        model:            Some("gpt-5.4".to_string()),
                        reasoning_effort: Some(ReasoningEffort::High),
                        speed:            Some(Speed::Fast),
                        permission_level: None,
                        capabilities:     vec![fabro_types::SessionCapability::Steer],
                    },
                    &scope,
                );
                Ok(CodergenResult::Text {
                    text:              "done".to_string(),
                    usage_by_model:    Vec::new(),
                    usage:             None,
                    files_touched:     Vec::new(),
                    last_file_touched: None,
                    timing:            StageTiming::default(),
                })
            }
        }

        let handler = AgentHandler::new(Some(Box::new(ProviderEventBackend)));
        let node = Node::new("step");
        let context = test_context();
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();
        let (services, run_store, logger) = make_services_with_run_store().await;

        handler
            .execute(&node, &context, &graph, tmp.path(), &services)
            .await
            .unwrap();
        logger.flush().await.unwrap();

        let state = run_store.state().await.unwrap();
        let node_state = state.stage(&StageId::new("step", 1)).unwrap();
        let provider_used = node_state.provider_used.as_ref().unwrap();
        assert_eq!(provider_used.provider.as_deref(), Some("openai"));
        assert_eq!(provider_used.reasoning_effort, Some(ReasoningEffort::High));
        assert_eq!(provider_used.speed, Some(Speed::Fast));
    }

    #[test]
    fn truncate_short_string() {
        assert_eq!(truncate("hello", 200), "hello");
    }

    #[test]
    fn truncate_long_string() {
        let long = "a".repeat(300);
        assert_eq!(truncate(&long, 200).len(), 200);
    }

    #[tokio::test]
    async fn codergen_handler_passes_thread_id_to_backend() {
        use std::sync::{Arc, Mutex};

        struct ThreadCapturingBackend {
            captured_thread_id: Arc<Mutex<Option<Option<String>>>>,
        }

        #[async_trait]
        impl CodergenBackend for ThreadCapturingBackend {
            async fn run(&self, request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
                *self.captured_thread_id.lock().unwrap() =
                    Some(request.thread_id.map(String::from));
                Ok(CodergenResult::Text {
                    text:              "ok".to_string(),
                    usage_by_model:    Vec::new(),
                    usage:             None,
                    files_touched:     Vec::new(),
                    last_file_touched: None,
                    timing:            StageTiming::default(),
                })
            }
        }

        let captured = Arc::new(Mutex::new(None));
        let backend = ThreadCapturingBackend {
            captured_thread_id: captured.clone(),
        };
        let handler = AgentHandler::new(Some(Box::new(backend)));

        let node = Node::new("work");
        let context = test_context();
        // Simulate what the engine stores in internal.thread_id
        context.set(keys::INTERNAL_THREAD_ID, serde_json::json!("main"));
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        handler
            .execute(&node, &context, &graph, tmp.path(), &make_services())
            .await
            .unwrap();

        let result = captured.lock().unwrap().clone();
        assert_eq!(result, Some(Some("main".to_string())));
    }

    #[tokio::test]
    async fn codergen_handler_passes_none_thread_id_when_absent() {
        use std::sync::{Arc, Mutex};

        struct ThreadCapturingBackend {
            captured_thread_id: Arc<Mutex<Option<Option<String>>>>,
        }

        #[async_trait]
        impl CodergenBackend for ThreadCapturingBackend {
            async fn run(&self, request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
                *self.captured_thread_id.lock().unwrap() =
                    Some(request.thread_id.map(String::from));
                Ok(CodergenResult::Text {
                    text:              "ok".to_string(),
                    usage_by_model:    Vec::new(),
                    usage:             None,
                    files_touched:     Vec::new(),
                    last_file_touched: None,
                    timing:            StageTiming::default(),
                })
            }
        }

        let captured = Arc::new(Mutex::new(None));
        let backend = ThreadCapturingBackend {
            captured_thread_id: captured.clone(),
        };
        let handler = AgentHandler::new(Some(Box::new(backend)));

        let node = Node::new("work");
        let context = test_context();
        // No thread context set
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        handler
            .execute(&node, &context, &graph, tmp.path(), &make_services())
            .await
            .unwrap();

        let result = captured.lock().unwrap().clone();
        assert_eq!(result, Some(None));
    }

    #[tokio::test]
    async fn codergen_handler_propagates_retryable_backend_error() {
        struct FailingBackend;

        #[async_trait]
        impl CodergenBackend for FailingBackend {
            async fn run(&self, _request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
                Err(Error::handler("Request timed out".to_string()))
            }
        }

        let handler = AgentHandler::new(Some(Box::new(FailingBackend)));
        let node = Node::new("step");
        let context = test_context();
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        let result = handler
            .execute(&node, &context, &graph, tmp.path(), &make_services())
            .await;
        let err = result.unwrap_err();
        assert!(err.is_retryable());
        assert!(err.to_string().contains("Request timed out"));
    }

    #[test]
    fn extract_status_fields_from_fenced_code_block() {
        let text = r#"Here is my analysis of the code.

```json
{"preferred_next_label": "fix", "outcome": "succeeded"}
```

That's it."#;
        let mut outcome = Outcome::success();
        extract_status_fields(text, &mut outcome);
        assert_eq!(outcome.preferred_label.as_deref(), Some("fix"));
    }

    #[test]
    fn extract_status_fields_from_bare_json() {
        let text = r#"I recommend routing to fix.
{"preferred_next_label": "fix_batch"}"#;
        let mut outcome = Outcome::success();
        extract_status_fields(text, &mut outcome);
        assert_eq!(outcome.preferred_label.as_deref(), Some("fix_batch"));
    }

    #[test]
    fn extract_status_fields_no_json() {
        let text = "Just some plain text response with no JSON at all.";
        let mut outcome = Outcome::success();
        extract_status_fields(text, &mut outcome);
        assert!(outcome.preferred_label.is_none());
        assert!(outcome.suggested_next_ids.is_empty());
    }

    #[test]
    fn extract_status_fields_json_without_status_fields() {
        let text = r#"Here is some data: {"name": "test", "count": 42}"#;
        let mut outcome = Outcome::success();
        extract_status_fields(text, &mut outcome);
        assert!(outcome.preferred_label.is_none());
        assert!(outcome.suggested_next_ids.is_empty());
    }

    #[test]
    fn extract_status_fields_context_updates_and_suggested_ids() {
        let text = r#"```json
{
  "preferred_next_label": "review",
  "suggested_next_ids": ["node_a", "node_b"],
  "context_updates": {"fix.files_changed": 3, "fix.summary": "patched"}
}
```"#;
        let mut outcome = Outcome::success();
        outcome
            .context_updates
            .insert("existing_key".to_string(), serde_json::json!("keep"));
        extract_status_fields(text, &mut outcome);
        assert_eq!(outcome.preferred_label.as_deref(), Some("review"));
        assert_eq!(outcome.suggested_next_ids, vec!["node_a", "node_b"]);
        assert_eq!(
            outcome.context_updates.get("fix.files_changed"),
            Some(&serde_json::json!(3))
        );
        assert_eq!(
            outcome.context_updates.get("fix.summary"),
            Some(&serde_json::json!("patched"))
        );
        // Existing keys preserved
        assert_eq!(
            outcome.context_updates.get("existing_key"),
            Some(&serde_json::json!("keep"))
        );
    }

    #[test]
    fn extract_status_fields_outcome_fail_with_reason() {
        let text = r#"{"outcome": "failed", "failure_reason": "tests failed"}"#;
        let mut outcome = Outcome::success();
        extract_status_fields(text, &mut outcome);
        assert_eq!(outcome.status, crate::outcome::StageOutcome::Failed {
            retry_requested: false,
        });
        assert_eq!(outcome.failure_reason(), Some("tests failed"));
    }

    #[test]
    fn extract_status_fields_outcome_success() {
        let text = r#"{"outcome": "succeeded"}"#;
        let mut outcome = Outcome::success();
        extract_status_fields(text, &mut outcome);
        assert_eq!(outcome.status, crate::outcome::StageOutcome::Succeeded);
        assert!(outcome.failure.is_none());
    }

    #[test]
    fn extract_status_fields_outcome_fail_without_reason() {
        let text = r#"{"outcome": "failed"}"#;
        let mut outcome = Outcome::success();
        extract_status_fields(text, &mut outcome);
        assert_eq!(outcome.status, crate::outcome::StageOutcome::Failed {
            retry_requested: false,
        });
        assert!(outcome.failure.is_none());
    }

    #[test]
    fn extract_status_fields_uses_last_match() {
        let text = r#"{"preferred_next_label": "first"}
Some text in between.
{"preferred_next_label": "second"}"#;
        let mut outcome = Outcome::success();
        extract_status_fields(text, &mut outcome);
        assert_eq!(outcome.preferred_label.as_deref(), Some("second"));
    }

    #[tokio::test]
    async fn codergen_handler_returns_fail_outcome_for_non_retryable_backend_error() {
        struct ValidationFailBackend;

        #[async_trait]
        impl CodergenBackend for ValidationFailBackend {
            async fn run(&self, _request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
                Err(Error::Validation("bad config".to_string()))
            }
        }

        let handler = AgentHandler::new(Some(Box::new(ValidationFailBackend)));
        let node = Node::new("step");
        let context = test_context();
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        let outcome = handler
            .execute(&node, &context, &graph, tmp.path(), &make_services())
            .await
            .unwrap();
        assert_eq!(outcome.status, crate::outcome::StageOutcome::Failed {
            retry_requested: false,
        });
        assert!(outcome.failure_reason().unwrap().contains("bad config"));
    }

    #[tokio::test]
    async fn codergen_handler_prepends_preamble_to_prompt() {
        use std::sync::{Arc, Mutex};

        struct PromptCapturingBackend {
            captured_prompt: Arc<Mutex<Option<String>>>,
        }

        #[async_trait]
        impl CodergenBackend for PromptCapturingBackend {
            async fn run(&self, request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
                *self.captured_prompt.lock().unwrap() = Some(request.prompt.to_string());
                Ok(CodergenResult::Text {
                    text:              "ok".to_string(),
                    usage_by_model:    Vec::new(),
                    usage:             None,
                    files_touched:     Vec::new(),
                    last_file_touched: None,
                    timing:            StageTiming::default(),
                })
            }
        }

        let captured = Arc::new(Mutex::new(None));
        let backend = PromptCapturingBackend {
            captured_prompt: captured.clone(),
        };
        let handler = AgentHandler::new(Some(Box::new(backend)));

        let mut node = Node::new("report");
        node.attrs.insert(
            "prompt".to_string(),
            AttrValue::String("Summarize the results".to_string()),
        );
        let context = test_context();
        context.set(
            keys::CURRENT_PREAMBLE,
            serde_json::json!("## Test Output\n10 passed, 0 failed"),
        );
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        handler
            .execute(&node, &context, &graph, tmp.path(), &make_services())
            .await
            .unwrap();

        let prompt = captured.lock().unwrap().clone().unwrap();
        assert!(
            prompt.starts_with("## Test Output\n10 passed, 0 failed"),
            "prompt should start with preamble, got: {prompt}"
        );
        assert!(
            prompt.ends_with("Summarize the results"),
            "prompt should end with original prompt, got: {prompt}"
        );
        assert!(
            prompt.contains("\n\nSummarize"),
            "preamble and prompt should be separated by blank line"
        );
    }

    #[tokio::test]
    async fn codergen_handler_no_preamble_when_empty() {
        use std::sync::{Arc, Mutex};

        struct PromptCapturingBackend {
            captured_prompt: Arc<Mutex<Option<String>>>,
        }

        #[async_trait]
        impl CodergenBackend for PromptCapturingBackend {
            async fn run(&self, request: CodergenRunRequest<'_>) -> Result<CodergenResult, Error> {
                *self.captured_prompt.lock().unwrap() = Some(request.prompt.to_string());
                Ok(CodergenResult::Text {
                    text:              "ok".to_string(),
                    usage_by_model:    Vec::new(),
                    usage:             None,
                    files_touched:     Vec::new(),
                    last_file_touched: None,
                    timing:            StageTiming::default(),
                })
            }
        }

        let captured = Arc::new(Mutex::new(None));
        let backend = PromptCapturingBackend {
            captured_prompt: captured.clone(),
        };
        let handler = AgentHandler::new(Some(Box::new(backend)));

        let mut node = Node::new("report");
        node.attrs.insert(
            "prompt".to_string(),
            AttrValue::String("Summarize the results".to_string()),
        );
        let context = test_context();
        // No preamble set -- context.get_string returns ""
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();

        handler
            .execute(&node, &context, &graph, tmp.path(), &make_services())
            .await
            .unwrap();

        let prompt = captured.lock().unwrap().clone().unwrap();
        assert_eq!(prompt, "Summarize the results");
    }

    #[tokio::test]
    async fn codergen_handler_preamble_written_to_prompt_md() {
        let handler = AgentHandler::new(None);
        let mut node = Node::new("report");
        node.attrs.insert(
            "prompt".to_string(),
            AttrValue::String("Summarize".to_string()),
        );
        let context = test_context();
        context.set(
            keys::CURRENT_PREAMBLE,
            serde_json::json!("## Script Output\nAll tests passed"),
        );
        let graph = Graph::new("test");
        let tmp = TempDir::new().unwrap();
        let (services, run_store, logger) = make_services_with_run_store().await;

        handler
            .execute(&node, &context, &graph, tmp.path(), &services)
            .await
            .unwrap();
        logger.flush().await.unwrap();

        let state = run_store.state().await.unwrap();
        let node_state = state.stage(&StageId::new("report", 1)).unwrap();
        let prompt_content = node_state.prompt.as_deref().unwrap();
        assert!(
            prompt_content.contains("## Script Output\nAll tests passed"),
            "prompt.md should contain preamble"
        );
        assert!(
            prompt_content.contains("Summarize"),
            "prompt.md should contain original prompt"
        );
    }
}