ag-agent 0.12.6

//! Shared app-server prompt shaping helpers.

use std::path::Path;

use ag_protocol::ProtocolSchemaInstructionMode;

use crate::agent;
use crate::agent::InstructionDeliveryMode;
use crate::app_server::{AppServerError, AppServerTurnRequest};
use crate::channel::AgentRequestKind;
use crate::model::turn_prompt::{TurnPrompt, TurnPromptTextSource};

/// Reads the latest replay transcript, preferring the live source over the
/// queued snapshot.
///
/// The live source accumulates all transcript messages in real time,
/// including content from a turn that failed mid-stream. When available, it
/// provides a more complete transcript than the snapshot captured at
/// turn-enqueue time.
pub(crate) fn read_latest_replay_transcript(request: &AppServerTurnRequest) -> Option<String> {
    if let Some(live_transcript) = &request.live_transcript
        && let Some(transcript_text) = live_transcript.replay_text()
        && !transcript_text.trim().is_empty()
    {
        return Some(transcript_text);
    }

    request.replay_transcript.clone()
}

/// Returns the turn prompt, applying protocol preamble and optional context
/// replay according to the selected instruction delivery mode.
///
///
/// `BootstrapFull` and `BootstrapWithReplay` include the shared protocol
/// preamble, while `DeltaOnly` emits only a compact reminder for provider
/// contexts that already received that contract. Providers that enforce the
/// response schema at transport level can request a policy-only bootstrap
/// preamble to avoid duplicating the full JSON Schema in prompt text.
///
/// # Errors
/// Returns an error when Askama prompt rendering fails after a context reset.
pub(crate) fn turn_prompt_for_runtime(
    prompt: impl Into<TurnPrompt>,
    request_kind: &AgentRequestKind,
    replay_transcript: Option<&str>,
    instruction_delivery_mode: InstructionDeliveryMode,
    schema_instruction_mode: ProtocolSchemaInstructionMode,
    workspace_root: &Path,
) -> Result<TurnPrompt, AppServerError> {
    let prompt = prompt.into();
    let agent_prompt = prompt.agent_text();
    let turn_prompt = agent::prepare_prompt_text(agent::PromptPreparationRequest {
        instruction_delivery_mode,
        prompt: &agent_prompt,
        protocol_profile: request_kind.protocol_profile(),
        replay_transcript,
        schema_instruction_mode,
        workspace_root,
    })
    .map_err(|error| AppServerError::PromptRender(error.to_string()))?;

    Ok(TurnPrompt {
        attachments: prompt.attachments,
        text: turn_prompt,
        text_source: TurnPromptTextSource::AgentData,
    })
}

/// Plans how one app-server turn should deliver Agentty's instruction
/// contract for the active runtime context.
pub(crate) fn instruction_delivery_mode_for_runtime(
    request: &AppServerTurnRequest,
    runtime_provider_conversation_id: Option<&str>,
    should_replay_transcript: bool,
) -> InstructionDeliveryMode {
    agent::plan_app_server_instruction_delivery(
        &request.request_kind,
        runtime_provider_conversation_id,
        request.persisted_instruction_conversation_id.as_deref(),
        should_replay_transcript,
    )
}

#[cfg(test)]
mod tests {
    use std::path::PathBuf;
    use std::sync::Arc;

    use super::*;
    use crate::channel::LiveTranscript;
    use crate::model::agent::ReasoningLevel;

    /// Workspace root used by app-server prompt shaping tests.
    const TEST_WORKSPACE_ROOT: &str = "/tmp/agentty-wt/session-1";

    #[derive(Debug)]
    struct TestLiveTranscript {
        text: String,
    }

    impl LiveTranscript for TestLiveTranscript {
        fn replay_text(&self) -> Option<String> {
            Some(self.text.clone())
        }
    }

    fn live_transcript(text: &str) -> Arc<dyn LiveTranscript> {
        Arc::new(TestLiveTranscript {
            text: text.to_string(),
        })
    }

    /// Returns one persisted bootstrap marker that matches the active
    /// app-server instruction contract for session turns.
    fn persisted_instruction_conversation_id_for_session_turn(
        provider_conversation_id: Option<&str>,
    ) -> Option<String> {
        agent::normalize_instruction_conversation_id(provider_conversation_id)
    }

    #[test]
    fn read_latest_replay_transcript_prefers_live_source() {
        // Arrange
        let request = AppServerTurnRequest {
            folder: PathBuf::from("/tmp/test"),
            live_transcript: Some(live_transcript("live content")),
            main_checkout_root: None,
            model: "test-model".to_string(),
            prompt: TurnPrompt::from("hello"),
            provider_conversation_id: None,
            persisted_instruction_conversation_id: None,
            reasoning_level: ReasoningLevel::default(),
            request_kind: AgentRequestKind::SessionStart,
            replay_transcript: None,
            session_id: "test-session".to_string(),
        };

        // Act
        let output = read_latest_replay_transcript(&request);

        // Assert
        assert_eq!(output, Some("live content".to_string()));
    }

    #[test]
    fn read_latest_replay_transcript_falls_back_when_live_source_is_empty() {
        // Arrange
        let request = AppServerTurnRequest {
            folder: PathBuf::from("/tmp/test"),
            live_transcript: Some(live_transcript("  ")),
            main_checkout_root: None,
            model: "test-model".to_string(),
            prompt: TurnPrompt::from("hello"),
            provider_conversation_id: None,
            persisted_instruction_conversation_id: None,
            reasoning_level: ReasoningLevel::default(),
            request_kind: AgentRequestKind::SessionStart,
            replay_transcript: Some("queued transcript".to_string()),
            session_id: "test-session".to_string(),
        };

        // Act
        let output = read_latest_replay_transcript(&request);

        // Assert
        assert_eq!(output, Some("queued transcript".to_string()));
    }

    #[test]
    fn read_latest_replay_transcript_returns_none_when_no_replay_text() {
        // Arrange
        let request = AppServerTurnRequest {
            folder: PathBuf::from("/tmp/test"),
            live_transcript: None,
            main_checkout_root: None,
            model: "test-model".to_string(),
            prompt: TurnPrompt::from("hello"),
            provider_conversation_id: None,
            persisted_instruction_conversation_id: None,
            reasoning_level: ReasoningLevel::default(),
            request_kind: AgentRequestKind::SessionStart,
            replay_transcript: None,
            session_id: "test-session".to_string(),
        };

        // Act
        let output = read_latest_replay_transcript(&request);

        // Assert
        assert!(output.is_none());
    }

    #[test]
    fn turn_prompt_for_runtime_includes_protocol_preamble() {
        // Arrange
        let prompt = TurnPrompt::from("fix the bug");
        let request_kind = AgentRequestKind::SessionStart;

        // Act
        let result = turn_prompt_for_runtime(
            prompt,
            &request_kind,
            None,
            InstructionDeliveryMode::BootstrapFull,
            ProtocolSchemaInstructionMode::PromptSchema,
            Path::new(TEST_WORKSPACE_ROOT),
        );

        // Assert
        let turn_prompt = result.expect("prompt rendering should succeed");
        assert!(turn_prompt.text.contains("fix the bug"));
        assert!(turn_prompt.text.contains("Structured response protocol:"));
        assert!(turn_prompt.text.contains("Anything outside that"));
        assert!(turn_prompt.text.contains("root is read-only."));
    }

    #[test]
    fn turn_prompt_for_runtime_omits_full_schema_for_transport_schema_mode() {
        // Arrange
        let prompt = TurnPrompt::from("fix the bug");
        let request_kind = AgentRequestKind::SessionStart;

        // Act
        let result = turn_prompt_for_runtime(
            prompt,
            &request_kind,
            None,
            InstructionDeliveryMode::BootstrapFull,
            ProtocolSchemaInstructionMode::TransportSchema,
            Path::new(TEST_WORKSPACE_ROOT),
        );

        // Assert
        let turn_prompt = result.expect("prompt rendering should succeed");
        assert!(turn_prompt.text.contains("Structured response protocol:"));
        assert!(
            turn_prompt
                .text
                .contains("provider enforces Agentty's response JSON schema")
        );
        assert!(!turn_prompt.text.contains("Authoritative JSON Schema:"));
    }

    #[test]
    fn turn_prompt_for_runtime_uses_compact_refresh_reminder_for_delta_only() {
        // Arrange
        let prompt = TurnPrompt::from("continue the fix");
        let request_kind = AgentRequestKind::SessionResume;

        // Act
        let result = turn_prompt_for_runtime(
            prompt,
            &request_kind,
            None,
            InstructionDeliveryMode::DeltaOnly,
            ProtocolSchemaInstructionMode::PromptSchema,
            Path::new(TEST_WORKSPACE_ROOT),
        );

        // Assert
        let turn_prompt = result.expect("prompt rendering should succeed");
        assert!(turn_prompt.text.contains("Protocol refresh reminder:"));
        assert!(!turn_prompt.text.contains("Authoritative JSON Schema:"));
    }

    #[test]
    fn turn_prompt_for_runtime_rewrites_user_at_lookups_for_agent_delivery() {
        // Arrange
        let prompt = TurnPrompt::from("review @src/main.rs");
        let request_kind = AgentRequestKind::SessionStart;

        // Act
        let result = turn_prompt_for_runtime(
            prompt,
            &request_kind,
            None,
            InstructionDeliveryMode::BootstrapFull,
            ProtocolSchemaInstructionMode::PromptSchema,
            Path::new(TEST_WORKSPACE_ROOT),
        );

        // Assert
        let turn_prompt = result.expect("prompt rendering should succeed");
        assert!(turn_prompt.text.contains("\"src/main.rs\""));
        assert!(!turn_prompt.text.contains("@src/main.rs"));
        assert!(!turn_prompt.text.contains("looked/up/"));
    }

    #[test]
    fn turn_prompt_for_runtime_preserves_generated_at_tokens_for_agent_data() {
        // Arrange
        let prompt = TurnPrompt::from_agent_data(
            "Review this diff:\n```diff\n+@dataclass\n+class Config:\n+    pass\n```".to_string(),
        );
        let request_kind = AgentRequestKind::UtilityPrompt;

        // Act
        let result = turn_prompt_for_runtime(
            prompt,
            &request_kind,
            None,
            InstructionDeliveryMode::BootstrapFull,
            ProtocolSchemaInstructionMode::PromptSchema,
            Path::new(TEST_WORKSPACE_ROOT),
        );

        // Assert
        let turn_prompt = result.expect("prompt rendering should succeed");
        assert!(turn_prompt.text.contains("+@dataclass"));
        assert!(!turn_prompt.text.contains("+\"dataclass\""));
    }

    #[test]
    fn instruction_delivery_mode_for_runtime_reuses_matching_bootstrap_state() {
        // Arrange
        let request = AppServerTurnRequest {
            folder: PathBuf::from("/tmp/test"),
            live_transcript: None,
            main_checkout_root: None,
            model: "test-model".to_string(),
            prompt: TurnPrompt::from("hello"),
            provider_conversation_id: Some("thread-123".to_string()),
            persisted_instruction_conversation_id:
                persisted_instruction_conversation_id_for_session_turn(Some("thread-123")),
            reasoning_level: ReasoningLevel::default(),
            request_kind: AgentRequestKind::SessionResume,
            replay_transcript: None,
            session_id: "test-session".to_string(),
        };

        // Act
        let delivery_mode =
            instruction_delivery_mode_for_runtime(&request, Some("thread-123"), false);

        // Assert
        assert_eq!(delivery_mode, InstructionDeliveryMode::DeltaOnly);
    }
}