ag_agent/agent/

provider.rs

//! Shared provider registry and transport policy descriptors.

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

#[cfg(test)]
use ag_protocol::AgentResponseSummary;
use ag_protocol::{
    AgentResponse, ProtocolRequestProfile, ProtocolSchemaInstructionMode,
    format_protocol_parse_debug_details, normalize_turn_response, parse_agent_response_strict,
};

use super::backend::{
    AgentBackend, AgentBackendError, AgentPromptTransport, AgentTransport, AppServerThoughtPolicy,
    BuildCommandRequest,
};
use super::prompt;
use super::response_parser::ParsedResponse;
use crate::app_server::AppServerClient;
use crate::model::agent::{AgentKind, AgentModel};

/// Factory hook used to build or override provider-specific app-server
/// clients.
type AppServerClientFactory =
    fn(Option<Arc<dyn AppServerClient>>) -> Option<Arc<dyn AppServerClient>>;

/// Creates the backend implementation for the selected agent provider.
pub fn create_backend(kind: AgentKind) -> Box<dyn AgentBackend> {
    (provider_descriptor(kind).backend_factory)()
}

/// Returns the app-server client for the selected provider when applicable.
pub fn create_app_server_client(
    kind: AgentKind,
    default_client: Option<Arc<dyn AppServerClient>>,
) -> Option<Arc<dyn AppServerClient>> {
    (provider_descriptor(kind).app_server_client_factory)(default_client)
}

/// Parses provider output and returns final response content and usage stats.
pub fn parse_response(kind: AgentKind, stdout: &str, stderr: &str) -> ParsedResponse {
    (provider_descriptor(kind).parse_response)(stdout, stderr)
}

/// Removes provider-owned worktree artifacts that are derived from one session
/// folder.
///
/// This is used by session teardown paths after provider setup or command
/// construction has created auxiliary filesystem state outside the real
/// worktree. Providers that do not create such state are no-ops.
///
/// # Errors
/// Returns an error when a provider-owned artifact exists but cannot be
/// removed safely.
pub fn cleanup_session_worktree_artifacts(folder: &Path) -> Result<(), AgentBackendError> {
    super::antigravity::cleanup_workspace_alias(folder)
}

/// Parses one stream line into incremental text and content classification.
///
/// Returns `(text, is_response_content)` where `is_response_content` is `true`
/// for model-authored content and `false` for progress updates.
pub fn parse_stream_output_line(kind: AgentKind, stdout_line: &str) -> Option<(String, bool)> {
    (provider_descriptor(kind).parse_stream_output_line)(stdout_line)
}

/// Returns transport mode for the selected provider.
pub fn transport_mode(kind: AgentKind) -> AgentTransport {
    provider_descriptor(kind).transport
}

/// Returns whether the provider expects prompts through stdin.
pub(crate) fn prompt_transport(kind: AgentKind) -> AgentPromptTransport {
    provider_descriptor(kind).prompt_transport
}

/// Returns whether bootstrap prompts should include schema text for the
/// selected provider.
///
/// Providers that enforce Agentty's response shape natively still receive
/// policy and field-routing instructions, but skip the large prompt-side JSON
/// Schema to avoid redundant tokens.
pub fn protocol_schema_instruction_mode(kind: AgentKind) -> ProtocolSchemaInstructionMode {
    provider_descriptor(kind).protocol_schema_instruction_mode
}

/// Parses one final assistant payload strictly against the shared protocol and
/// normalizes it for the active request profile.
///
/// # Errors
/// Returns a descriptive error when provider output does not match the
/// required protocol JSON, including parse diagnostics that help explain why
/// the payload was rejected.
pub fn parse_turn_response(
    kind: AgentKind,
    response_text: &str,
    protocol_profile: ProtocolRequestProfile,
) -> Result<AgentResponse, String> {
    let response = parse_agent_response_strict(response_text).map_err(|error| {
        format!(
            "Agent output did not match the required JSON schema from {kind}: \
             {error}\nprotocol_profile: \
             {protocol_profile:?}\ndebug_details:\n{}\nresponse:\n{response_text}",
            format_protocol_parse_debug_details(response_text)
        )
    })?;

    Ok(normalize_turn_response(response, protocol_profile))
}

/// Returns whether one app-server assistant chunk should be treated as
/// thought text instead of transcript output.
pub fn is_app_server_thought_chunk(kind: AgentKind, is_delta: bool, phase: Option<&str>) -> bool {
    if !is_delta {
        return false;
    }

    match provider_descriptor(kind).app_server_thought_policy {
        AppServerThoughtPolicy::None => false,
        AppServerThoughtPolicy::PhaseLabel => phase.is_some_and(is_codex_thought_phase_label),
    }
}

/// Builds one optional stdin payload for providers that stream prompts instead
/// of sending them through argv.
///
/// # Errors
/// Returns an error when provider-specific prompt rendering fails.
pub fn build_command_stdin_payload(
    kind: AgentKind,
    request: BuildCommandRequest<'_>,
) -> Result<Option<Vec<u8>>, AgentBackendError> {
    let protocol_schema_instruction_mode = protocol_schema_instruction_mode(kind);

    match prompt_transport(kind) {
        AgentPromptTransport::Argv => Ok(None),
        AgentPromptTransport::Stdin => match kind {
            AgentKind::Antigravity => prompt::build_prompt_stdin_payload(
                request,
                protocol_schema_instruction_mode,
                "Antigravity",
            )
            .map(Some),
            AgentKind::Claude => prompt::build_prompt_stdin_payload(
                request,
                protocol_schema_instruction_mode,
                "Claude",
            )
            .map(Some),
            AgentKind::Codex | AgentKind::Gemini => Ok(None),
        },
    }
}

/// Parses one app-server model string into its routing provider kind.
///
/// # Errors
/// Returns an error when `model` is not a known app-server-backed
/// [`AgentModel`].
pub fn provider_kind_for_model(model: &str) -> Result<AgentKind, String> {
    let model = model
        .parse::<AgentModel>()
        .map_err(|error| format!("unknown model `{model}`: {error}"))?;
    if AgentKind::Codex.supports_model(model) {
        return Ok(AgentKind::Codex);
    }

    if AgentKind::Gemini.supports_model(model) {
        return Ok(AgentKind::Gemini);
    }

    Err(format!(
        "model `{}` does not use app-server routing",
        model.as_str()
    ))
}

/// One backend/provider descriptor containing construction and parsing hooks.
struct AgentProviderDescriptor {
    app_server_client_factory: AppServerClientFactory,
    app_server_thought_policy: AppServerThoughtPolicy,
    backend_factory: fn() -> Box<dyn AgentBackend>,
    parse_response: fn(&str, &str) -> ParsedResponse,
    parse_stream_output_line: fn(&str) -> Option<(String, bool)>,
    prompt_transport: AgentPromptTransport,
    protocol_schema_instruction_mode: ProtocolSchemaInstructionMode,
    transport: AgentTransport,
}

fn provider_descriptor(kind: AgentKind) -> AgentProviderDescriptor {
    match kind {
        AgentKind::Antigravity => AgentProviderDescriptor {
            app_server_client_factory: |_default_client| None,
            app_server_thought_policy: AppServerThoughtPolicy::None,
            backend_factory: || Box::new(super::antigravity::AntigravityBackend),
            parse_response: super::response_parser::parse_antigravity_response_with_fallback,
            parse_stream_output_line: super::response_parser::parse_antigravity_stream_output_line,
            prompt_transport: AgentPromptTransport::Stdin,
            protocol_schema_instruction_mode: ProtocolSchemaInstructionMode::PromptSchema,
            transport: AgentTransport::Cli,
        },
        AgentKind::Gemini => AgentProviderDescriptor {
            app_server_client_factory: |default_client| {
                Some(default_client.unwrap_or_else(|| {
                    Arc::new(super::app_server::RealGeminiAcpClient::new())
                        as Arc<dyn AppServerClient>
                }))
            },
            app_server_thought_policy: AppServerThoughtPolicy::None,
            backend_factory: || Box::new(super::gemini::GeminiBackend),
            parse_response: super::response_parser::parse_gemini_response_with_fallback,
            parse_stream_output_line: super::response_parser::parse_gemini_stream_output_line,
            prompt_transport: AgentPromptTransport::Argv,
            protocol_schema_instruction_mode: ProtocolSchemaInstructionMode::PromptSchema,
            transport: AgentTransport::AppServer,
        },
        AgentKind::Claude => AgentProviderDescriptor {
            app_server_client_factory: |_default_client| None,
            app_server_thought_policy: AppServerThoughtPolicy::None,
            backend_factory: || Box::new(super::claude::ClaudeBackend),
            parse_response: super::response_parser::parse_claude_response_with_fallback,
            parse_stream_output_line: super::response_parser::parse_claude_stream_output_line,
            prompt_transport: AgentPromptTransport::Stdin,
            protocol_schema_instruction_mode: ProtocolSchemaInstructionMode::TransportSchema,
            transport: AgentTransport::Cli,
        },
        AgentKind::Codex => AgentProviderDescriptor {
            app_server_client_factory: |default_client| {
                Some(default_client.unwrap_or_else(|| {
                    Arc::new(super::app_server::RealCodexAppServerClient::new())
                        as Arc<dyn AppServerClient>
                }))
            },
            app_server_thought_policy: AppServerThoughtPolicy::PhaseLabel,
            backend_factory: || Box::new(super::codex::CodexBackend),
            parse_response: super::response_parser::parse_codex_response_with_fallback,
            parse_stream_output_line: super::response_parser::parse_codex_stream_output_line,
            prompt_transport: AgentPromptTransport::Argv,
            protocol_schema_instruction_mode: ProtocolSchemaInstructionMode::TransportSchema,
            transport: AgentTransport::AppServer,
        },
    }
}

/// Returns whether one Codex phase label denotes thought/planning text.
///
/// Phase matching is case-insensitive so provider variants such as `Thinking`
/// and `PLAN` continue to route to thought deltas.
fn is_codex_thought_phase_label(phase: &str) -> bool {
    let normalized_phase = phase.trim();

    normalized_phase.eq_ignore_ascii_case("thinking")
        || normalized_phase.eq_ignore_ascii_case("plan")
        || normalized_phase.eq_ignore_ascii_case("reasoning")
        || normalized_phase.eq_ignore_ascii_case("thought")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    /// Ensures transport capability is provided by infra backend descriptors,
    /// not domain enums.
    fn test_transport_mode_reports_expected_transport_by_provider() {
        // Arrange
        let antigravity_kind = AgentKind::Antigravity;
        let claude_kind = AgentKind::Claude;
        let codex_kind = AgentKind::Codex;
        let gemini_kind = AgentKind::Gemini;

        // Act
        let antigravity_transport = transport_mode(antigravity_kind);
        let claude_transport = transport_mode(claude_kind);
        let codex_transport = transport_mode(codex_kind);
        let gemini_transport = transport_mode(gemini_kind);

        // Assert
        assert_eq!(antigravity_transport, AgentTransport::Cli);
        assert_eq!(claude_transport, AgentTransport::Cli);
        assert_eq!(codex_transport, AgentTransport::AppServer);
        assert_eq!(gemini_transport, AgentTransport::AppServer);
    }

    #[test]
    /// Ensures prompt delivery is also derived from the shared provider
    /// descriptor.
    fn test_prompt_transport_reports_expected_mode_by_provider() {
        // Arrange
        let antigravity_kind = AgentKind::Antigravity;
        let claude_kind = AgentKind::Claude;
        let codex_kind = AgentKind::Codex;
        let gemini_kind = AgentKind::Gemini;

        // Act
        let antigravity_transport = prompt_transport(antigravity_kind);
        let claude_transport = prompt_transport(claude_kind);
        let codex_transport = prompt_transport(codex_kind);
        let gemini_transport = prompt_transport(gemini_kind);

        // Assert
        assert_eq!(antigravity_transport, AgentPromptTransport::Stdin);
        assert_eq!(claude_transport, AgentPromptTransport::Stdin);
        assert_eq!(codex_transport, AgentPromptTransport::Argv);
        assert_eq!(gemini_transport, AgentPromptTransport::Argv);
    }

    #[test]
    /// Ensures provider schema capabilities are derived from the shared
    /// descriptor.
    fn test_protocol_schema_instruction_mode_reports_expected_mode_by_provider() {
        // Arrange / Act / Assert
        assert_eq!(
            protocol_schema_instruction_mode(AgentKind::Antigravity),
            ProtocolSchemaInstructionMode::PromptSchema
        );
        assert_eq!(
            protocol_schema_instruction_mode(AgentKind::Gemini),
            ProtocolSchemaInstructionMode::PromptSchema
        );
        assert_eq!(
            protocol_schema_instruction_mode(AgentKind::Claude),
            ProtocolSchemaInstructionMode::TransportSchema
        );
        assert_eq!(
            protocol_schema_instruction_mode(AgentKind::Codex),
            ProtocolSchemaInstructionMode::TransportSchema
        );
    }

    #[test]
    /// Ensures providers reject malformed final protocol payloads.
    fn test_parse_turn_response_rejects_invalid_payload() {
        // Arrange
        let raw_response = "plain response";

        for kind in [
            AgentKind::Antigravity,
            AgentKind::Claude,
            AgentKind::Codex,
            AgentKind::Gemini,
        ] {
            // Act
            let error =
                parse_turn_response(kind, raw_response, ProtocolRequestProfile::SessionTurn)
                    .expect_err("plain response should fail strict protocol parsing");

            // Assert
            assert!(error.contains("debug_details:"));
            assert!(error.contains("first_non_whitespace_char: 'p'"));
            assert!(error.contains("direct_json_error_location: line 1, column 1"));
        }
    }

    #[test]
    /// Ensures valid session-turn payloads still gain an empty summary when
    /// the response omits it.
    fn test_parse_turn_response_fills_missing_summary_for_session_turn() {
        // Arrange
        let raw_response = r#"{"answer":"done","questions":[],"summary":null}"#;

        // Act
        let result = parse_turn_response(
            AgentKind::Codex,
            raw_response,
            ProtocolRequestProfile::SessionTurn,
        )
        .expect("valid protocol response should parse");

        // Assert
        assert_eq!(result.answer, "done");
        assert_eq!(
            result.summary,
            Some(AgentResponseSummary {
                session: String::new(),
                turn: String::new(),
            })
        );
    }

    #[test]
    /// Ensures Codex app-server phase labels map to thought deltas through the
    /// shared provider descriptor.
    fn test_is_app_server_thought_chunk_reports_codex_phase_labels() {
        // Arrange / Act / Assert
        assert!(is_app_server_thought_chunk(
            AgentKind::Codex,
            true,
            Some("thinking"),
        ));
        assert!(!is_app_server_thought_chunk(
            AgentKind::Gemini,
            true,
            Some("thinking"),
        ));
    }

    #[test]
    /// Ensures model strings resolve through the shared provider registry.
    fn test_provider_kind_for_model_reports_owner() {
        // Arrange / Act / Assert
        assert_eq!(
            provider_kind_for_model(AgentModel::Gemini31ProPreview.as_str()).expect("known model"),
            AgentKind::Gemini
        );
        assert_eq!(
            provider_kind_for_model(AgentModel::Gpt55.as_str()).expect("known model"),
            AgentKind::Codex
        );
    }
}