code-muse 0.0.1__py3-none-any.whl

code_muse/plugins/azure_foundry/register_callbacks.py

@@ -0,0 +1,495 @@
+"""Azure AI Foundry Plugin callbacks for Muse CLI.
+
+This plugin enables Muse to use Anthropic Claude models hosted on
+Microsoft Azure AI Foundry with Azure AD (Entra ID) authentication.
+
+The plugin uses credentials from `az login` to authenticate, eliminating
+the need for API keys.
+"""
+
+import logging
+import sys
+from typing import Any
+
+from code_muse.callbacks import register_callback
+from code_muse.command_line.utils import safe_input
+from code_muse.messaging import emit_error, emit_info, emit_success, emit_warning
+from code_muse.tools.command_runner import set_awaiting_user_input
+
+from .config import (
+    DEFAULT_DEPLOYMENT_NAMES,
+    ENV_FOUNDRY_RESOURCE,
+    get_foundry_resource,
+)
+from .discovery import find_account, list_deployments
+from .token import get_token_provider
+from .utils import (
+    add_discovered_models_to_config,
+    add_foundry_models_to_config,
+    get_foundry_models_from_config,
+    remove_foundry_models_from_config,
+    resolve_env_var,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# ============================================================================
+# Slash Command Handlers
+# ============================================================================
+
+
+def _handle_foundry_status() -> None:
+    """Handle the /foundry-status command.
+
+    Displays the current Azure AD authentication status and configured
+    Foundry models.
+    """
+    emit_info("")
+    emit_info("Azure AI Foundry Status")
+    emit_info("=" * 40)
+
+    # Check Azure AD authentication
+    token_provider = get_token_provider()
+    is_auth, status_msg, user_info = token_provider.check_auth_status()
+
+    if is_auth:
+        emit_success(f"Authentication: {status_msg}")
+        if user_info:
+            emit_info(f"   Logged in as: {user_info}")
+    else:
+        emit_warning(f"Authentication: {status_msg}")
+
+    # List configured models and check resource
+    foundry_models = get_foundry_models_from_config()
+
+    # Check resource - from env var or from configured models
+    resource = get_foundry_resource()
+    if not resource and foundry_models:
+        # Get resource from first configured model
+        first_model = next(iter(foundry_models.values()))
+        resource = first_model.get("foundry_resource", "")
+        if resource.startswith("$"):
+            resource = None  # It's an unresolved env var reference
+
+    emit_info("")
+    if resource:
+        emit_info(f"Foundry Resource: {resource}")
+    else:
+        emit_warning(f"Foundry Resource: Not set (set {ENV_FOUNDRY_RESOURCE})")
+
+    emit_info("")
+    if foundry_models:
+        emit_info(f"Configured Models ({len(foundry_models)}):")
+        for model_key, config in foundry_models.items():
+            deployment = config.get("name", "unknown")
+            emit_info(f"   - {model_key}: {deployment}")
+    else:
+        emit_info("Configured Models: None")
+        emit_info("   Run /foundry-setup to configure models")
+
+    emit_info("")
+
+
+def _handle_foundry_setup() -> None:
+    """Handle the /foundry-setup command.
+
+    Interactive wizard to configure Azure Foundry models.
+    Uses print() for synchronous output to avoid message bus buffering issues.
+    """
+
+    def _print(msg: str = "") -> None:
+        """Print with immediate flush."""
+        print(msg, flush=True)
+
+    _print()
+    _print("Azure AI Foundry Setup")
+    _print("=" * 40)
+    _print()
+
+    # Check Azure CLI authentication first
+    _print("Step 1: Checking Azure CLI authentication...")
+    token_provider = get_token_provider()
+    is_auth, status_msg, user_info = token_provider.check_auth_status()
+
+    if not is_auth:
+        _print(f"   ERROR: {status_msg}")
+        _print()
+        _print("Please run 'az login' first, then try again.")
+        return
+
+    _print(f"   OK: {status_msg}")
+    if user_info:
+        _print(f"   User: {user_info}")
+    _print()
+
+    # Get resource name
+    _print("Step 2: Azure Resource Name")
+    current_resource = get_foundry_resource()
+    if current_resource:
+        _print(f"   Current: {current_resource}")
+
+    resource_prompt = "   Enter resource name"
+    if current_resource:
+        resource_prompt += f" [{current_resource}]"
+    resource_prompt += ": "
+
+    set_awaiting_user_input(True)
+    try:
+        sys.stdout.flush()
+        resource_input = safe_input(resource_prompt).strip()
+        resource_name = resource_input if resource_input else current_resource
+
+        if not resource_name:
+            _print("   ERROR: Resource name is required.")
+            return
+
+        _print()
+
+        # Step 3: Try auto-discovery, fall back to manual
+        _print("Step 3: Discovering deployments...")
+        account = find_account(resource_name)
+
+        if account:
+            _print(f"   Found: {account.name} ({account.location})")
+            _print(f"   RG: {account.resource_group}")
+            _print()
+
+            deployments = list_deployments(account)
+            succeeded = [d for d in deployments if d.provisioning_state == "Succeeded"]
+
+            if succeeded:
+                _print(f"   {len(succeeded)} active deployment(s):")
+                for d in succeeded:
+                    _print(f"   - {d.name} ({d.model_format}: {d.model_name})")
+                _print()
+
+                sys.stdout.flush()
+                confirm = safe_input("   Configure these? [Y/n]: ").strip().lower()
+                if confirm not in ("", "y", "yes"):
+                    _print("   Skipped.")
+                    return
+            else:
+                _print("   No active deployments found.")
+                return
+        else:
+            _print("   Discovery failed — falling back to manual entry.")
+            _print()
+            succeeded = None
+
+    except KeyboardInterrupt, EOFError:
+        _print()
+        _print("Setup cancelled.")
+        return
+    finally:
+        set_awaiting_user_input(False)
+
+    _print()
+
+    # Step 4: Save configuration
+    _print("Step 4: Saving configuration...")
+
+    if not get_foundry_resource():
+        _print(
+            f"   Tip: Set {ENV_FOUNDRY_RESOURCE}={resource_name} in your environment"
+        )
+
+    if succeeded is not None:
+        # Auto-discovered — configure all succeeded deployments
+        added_models = add_discovered_models_to_config(resource_name, succeeded)
+    else:
+        # Manual fallback — use hardcoded Anthropic defaults
+        added_models = add_foundry_models_to_config(
+            resource_name=resource_name,
+            opus_deployment=DEFAULT_DEPLOYMENT_NAMES["opus"],
+            sonnet_deployment=DEFAULT_DEPLOYMENT_NAMES["sonnet"],
+            haiku_deployment=DEFAULT_DEPLOYMENT_NAMES["haiku"],
+        )
+
+    _print()
+    if added_models:
+        _print(f"OK: Configured {len(added_models)} model(s):")
+        for model_key in added_models:
+            _print(f"   - {model_key}")
+        _print()
+        _print(f"Use '/model {added_models[0]}' to switch to a Foundry model.")
+    else:
+        _print("WARNING: No models were added. Check the configuration.")
+
+    _print()
+
+
+def _handle_foundry_remove() -> None:
+    """Handle the /foundry-remove command.
+
+    Removes all Azure Foundry model configurations.
+    """
+    removed = remove_foundry_models_from_config()
+    if removed:
+        emit_success(f"Removed {len(removed)} Foundry model(s):")
+        for model_key in removed:
+            emit_info(f"   - {model_key}")
+    else:
+        emit_info("No Foundry models found in configuration.")
+
+
+# ============================================================================
+# Custom Command Registration
+# ============================================================================
+
+
+def _custom_help() -> list[tuple[str, str]]:
+    """Return help entries for custom commands."""
+    return [
+        (
+            "foundry-status",
+            "Check Azure AI Foundry authentication and configuration status",
+        ),
+        ("foundry-setup", "Interactive wizard to configure Azure Foundry models"),
+        ("foundry-remove", "Remove all Azure Foundry model configurations"),
+    ]
+
+
+def _handle_custom_command(command: str, name: str) -> bool | None:
+    """Handle custom slash commands for the Azure Foundry plugin.
+
+    Dispatches to handlers for "foundry-status", "foundry-setup", and
+    "foundry-remove" commands.
+
+    Args:
+        command: The full command string.
+        name: The command name (without slash).
+
+    Returns:
+        True if the command was handled successfully.
+        False if the command was recognized but handler() raised an exception.
+        None if the command is not handled by this plugin.
+    """
+    handlers = {
+        "foundry-status": _handle_foundry_status,
+        "foundry-setup": _handle_foundry_setup,
+        "foundry-remove": _handle_foundry_remove,
+    }
+
+    handler = handlers.get(name)
+    if handler is None:
+        return None
+
+    try:
+        handler()
+        return True
+    except Exception as e:
+        logger.exception("Error handling /%s command: %s", name, e)
+        emit_error(f"Command /{name} failed: {e}")
+        return False
+
+
+# ============================================================================
+# Model Type Handler
+# ============================================================================
+
+
+def _create_azure_foundry_model(
+    model_name: str, model_config: dict, config: dict
+) -> Any:
+    """Create an Azure Foundry model instance.
+
+    This handler is registered via the 'register_model_type' callback to handle
+    models with type='azure_foundry'.
+
+    Args:
+        model_name: The model key name (e.g., 'foundry-claude-opus').
+        model_config: The model configuration dictionary.
+        config: The full models configuration.
+
+    Returns:
+        An AnthropicModel instance configured for Azure Foundry, or None on error.
+    """
+    try:
+        from anthropic import AsyncAnthropicFoundry
+        from pydantic_ai.models.anthropic import AnthropicModel
+    except ImportError as e:
+        emit_error(
+            f"Failed to create Azure Foundry model '{model_name}': "
+            f"Missing dependency - {e}"
+        )
+        return None
+
+    from code_muse.claude_cache_client import patch_anthropic_client_messages
+    from code_muse.config import get_effective_model_settings
+    from code_muse.model_factory import CONTEXT_1M_BETA
+    from code_muse.provider_identity import (
+        make_anthropic_provider,
+        resolve_provider_identity,
+    )
+
+    # Get the Foundry resource name
+    resource_config = model_config.get("foundry_resource", f"${ENV_FOUNDRY_RESOURCE}")
+    resource_name = resolve_env_var(resource_config)
+
+    if not resource_name:
+        emit_warning(
+            f"Azure Foundry resource not configured for model '{model_name}'. "
+            f"Set {ENV_FOUNDRY_RESOURCE} or run /foundry-setup."
+        )
+        return None
+
+    # Get the deployment name (model name in Azure)
+    deployment_name = model_config.get("name")
+    if not deployment_name:
+        emit_warning(f"Deployment name not specified for model '{model_name}'.")
+        return None
+
+    # Get the token provider
+    token_provider = get_token_provider()
+
+    # Check authentication status
+    is_auth, status_msg, _ = token_provider.check_auth_status()
+    if not is_auth:
+        emit_warning(
+            f"Azure AD authentication failed for model '{model_name}': {status_msg}"
+        )
+        return None
+
+    try:
+        # Check for interleaved thinking setting (default True for Foundry models)
+        effective_settings = get_effective_model_settings(model_name)
+        interleaved_thinking = effective_settings.get("interleaved_thinking", True)
+
+        # Build anthropic-beta header if needed
+        beta_parts: list[str] = []
+        if interleaved_thinking:
+            beta_parts.append("interleaved-thinking-2025-05-14")
+
+        # Add 1M context beta header for long-context models
+        context_length = model_config.get("context_length", 200000)
+        if context_length >= 1_000_000:
+            beta_parts.append(CONTEXT_1M_BETA)
+
+        # Build default headers dict if we have beta features
+        default_headers: dict[str, str] | None = None
+        if beta_parts:
+            default_headers = {"anthropic-beta": ",".join(beta_parts)}
+
+        # Create the Azure Foundry Anthropic client with token provider
+        # Note: We pass default_headers here because AsyncAnthropicFoundry.with_options()
+        # has a bug where copy() passes auth_token which isn't a valid __init__ param
+        anthropic_client = AsyncAnthropicFoundry(
+            resource=resource_name,
+            azure_ad_token_provider=token_provider.get_token,
+            default_headers=default_headers,
+        )
+
+        # Patch for cache control injection
+        patch_anthropic_client_messages(anthropic_client)
+
+        # Create the pydantic-ai provider and model
+        provider_identity = resolve_provider_identity(model_name, model_config)
+        provider = make_anthropic_provider(
+            provider_identity,
+            anthropic_client=anthropic_client,
+        )
+
+        model = AnthropicModel(model_name=deployment_name, provider=provider)
+        logger.info(
+            "Created Azure Foundry model: %s -> %s @ %s",
+            model_name,
+            deployment_name,
+            resource_name,
+        )
+        return model
+
+    except Exception as e:
+        emit_error(f"Failed to create Azure Foundry model '{model_name}': {e}")
+        logger.exception(f"Error creating Azure Foundry model: {e}")
+        return None
+
+
+def _create_azure_foundry_openai_model(
+    model_name: str, model_config: dict, config: dict
+) -> Any:
+    """Create an Azure Foundry OpenAI model instance.
+
+    Handles models with type='azure_foundry_openai' — OpenAI models on
+    Azure AI Services using Azure AD token auth (no API keys).
+    """
+    try:
+        from openai import AsyncAzureOpenAI
+        from pydantic_ai.models.openai import OpenAIChatModel, OpenAIResponsesModel
+    except ImportError as e:
+        emit_error(f"Failed to create Azure Foundry OpenAI model '{model_name}': {e}")
+        return None
+
+    from code_muse.provider_identity import (
+        make_openai_provider,
+        resolve_provider_identity,
+    )
+
+    resource_config = model_config.get("foundry_resource", f"${ENV_FOUNDRY_RESOURCE}")
+    resource_name = resolve_env_var(resource_config)
+
+    if not resource_name:
+        emit_warning(
+            f"Azure Foundry resource not configured for model '{model_name}'. "
+            f"Set {ENV_FOUNDRY_RESOURCE} or run /foundry-setup."
+        )
+        return None
+
+    deployment_name = model_config.get("name")
+    if not deployment_name:
+        emit_warning(f"Deployment name not specified for model '{model_name}'.")
+        return None
+
+    token_provider = get_token_provider()
+    is_auth, status_msg, _ = token_provider.check_auth_status()
+    if not is_auth:
+        emit_warning(f"Azure AD auth failed for model '{model_name}': {status_msg}")
+        return None
+
+    try:
+        api_version = model_config.get("api_version", "2025-04-01-preview")
+        azure_endpoint = f"https://{resource_name}.openai.azure.com"
+
+        azure_client = AsyncAzureOpenAI(
+            azure_endpoint=azure_endpoint,
+            api_version=api_version,
+            azure_ad_token_provider=token_provider.get_token,
+        )
+
+        provider_identity = resolve_provider_identity(model_name, model_config)
+        provider = make_openai_provider(provider_identity, openai_client=azure_client)
+
+        if deployment_name.startswith("gpt-5"):
+            model = OpenAIResponsesModel(model_name=deployment_name, provider=provider)
+        else:
+            model = OpenAIChatModel(model_name=deployment_name, provider=provider)
+        logger.info(
+            "Created Azure Foundry OpenAI model: %s -> %s @ %s",
+            model_name,
+            deployment_name,
+            resource_name,
+        )
+        return model
+
+    except Exception as e:
+        emit_error(f"Failed to create Azure Foundry OpenAI model '{model_name}': {e}")
+        logger.exception("Error creating Azure Foundry OpenAI model: %s", e)
+        return None
+
+
+def _register_model_types() -> list[dict[str, Any]]:
+    """Register azure_foundry and azure_foundry_openai model type handlers."""
+    return [
+        {"type": "azure_foundry", "handler": _create_azure_foundry_model},
+        {"type": "azure_foundry_openai", "handler": _create_azure_foundry_openai_model},
+    ]
+
+
+# ============================================================================
+# Callback Registration
+# ============================================================================
+
+# Register all callbacks when this module is imported
+register_callback("custom_command_help", _custom_help)
+register_callback("custom_command", _handle_custom_command)
+register_callback("register_model_type", _register_model_types)

code_muse/plugins/azure_foundry/token.py

@@ -0,0 +1,180 @@
+"""Azure AD token provider for Azure AI Foundry authentication.
+
+This module provides token management for authenticating with Azure AI Foundry
+using credentials from the Azure CLI (`az login`).
+
+The token provider uses `AzureCliCredential` from the `azure-identity` library
+to obtain tokens without requiring API keys.
+"""
+
+import logging
+import time
+from collections.abc import Callable
+
+from .config import AZURE_COGNITIVE_SCOPE, TOKEN_REFRESH_BUFFER
+
+logger = logging.getLogger(__name__)
+
+# Singleton instance of the token provider
+_token_provider_instance: AzureFoundryTokenProvider | None = None
+
+
+class AzureFoundryTokenProvider:
+    """Provides Azure AD tokens for Anthropic Foundry using az login credentials.
+
+    This class wraps the Azure CLI credential to provide tokens for
+    authenticating with Azure AI Foundry. It handles token caching and
+    provides status checking functionality.
+
+    Example:
+        >>> provider = AzureFoundryTokenProvider()
+        >>> token = provider.get_token()
+        >>> # Use token for API calls
+    """
+
+    def __init__(self, scope: str = AZURE_COGNITIVE_SCOPE):
+        """Initialize the token provider.
+
+        Args:
+            scope: The Azure AD scope for token acquisition.
+                   Defaults to the Cognitive Services scope.
+        """
+        self._scope = scope
+        self._credential = None
+        self._token_provider_func: Callable[[], str] | None = None
+        self._initialized = False
+        self._init_error: str | None = None
+
+    def _ensure_initialized(self) -> bool:
+        """Lazily initialize the Azure credential.
+
+        Returns:
+            True if initialization succeeded, False otherwise.
+        """
+        if self._initialized:
+            return self._init_error is None
+
+        try:
+            from azure.identity import AzureCliCredential, get_bearer_token_provider
+
+            self._credential = AzureCliCredential()
+            self._token_provider_func = get_bearer_token_provider(
+                self._credential, self._scope
+            )
+            self._initialized = True
+            self._init_error = None
+            logger.debug("Azure CLI credential initialized successfully")
+            return True
+
+        except ImportError as e:
+            self._initialized = True
+            self._init_error = f"azure-identity package not installed: {e}"
+            logger.error(self._init_error)
+            return False
+
+        except Exception as e:
+            self._initialized = True
+            self._init_error = f"Failed to initialize Azure credential: {e}"
+            logger.error(self._init_error)
+            return False
+
+    def get_token(self) -> str:
+        """Get a valid access token for Azure AI Foundry.
+
+        The token is obtained from the Azure CLI credential and is
+        automatically refreshed when needed by the underlying provider.
+
+        Returns:
+            A valid access token string.
+
+        Raises:
+            RuntimeError: If the token provider is not initialized or
+                         if token acquisition fails.
+        """
+        if not self._ensure_initialized():
+            raise RuntimeError(self._init_error or "Token provider not initialized")
+
+        if self._token_provider_func is None:
+            raise RuntimeError("Token provider function not available")
+
+        try:
+            return self._token_provider_func()
+        except Exception as e:
+            logger.error(f"Failed to acquire token: {e}")
+            raise RuntimeError(f"Failed to acquire Azure AD token: {e}") from e
+
+    def check_auth_status(self) -> tuple[bool, str, str | None]:
+        """Check if Azure CLI authentication is valid.
+
+        Returns:
+            A tuple of (is_authenticated, status_message, user_info):
+            - is_authenticated: True if auth is valid, False otherwise
+            - status_message: Human-readable status description
+            - user_info: Email/UPN of logged-in user if available
+        """
+        if not self._ensure_initialized():
+            return False, self._init_error or "Not initialized", None
+
+        try:
+            # Try to get a token to verify authentication
+            token = self._credential.get_token(self._scope)
+            expires_in = int(token.expires_on - time.time())
+
+            # Try to get user info from the token
+            user_info = None
+            try:
+                # The token is a JWT, we can decode the payload to get user info
+                import base64
+                import json
+
+                # Split the JWT and decode the payload (second part)
+                parts = token.token.split(".")
+                if len(parts) >= 2:
+                    # Add padding if needed
+                    payload = parts[1]
+                    padding = 4 - len(payload) % 4
+                    if padding != 4:
+                        payload += "=" * padding
+                    decoded = base64.urlsafe_b64decode(payload)
+                    claims = json.loads(decoded)
+                    user_info = (
+                        claims.get("upn")
+                        or claims.get("email")
+                        or claims.get("preferred_username")
+                    )
+            except Exception:
+                # Ignore errors in user info extraction
+                pass
+
+            if expires_in > TOKEN_REFRESH_BUFFER:
+                return True, f"Valid (expires in {expires_in // 60} minutes)", user_info
+            else:
+                return (
+                    True,
+                    f"Valid but expiring soon ({expires_in} seconds)",
+                    user_info,
+                )
+
+        except Exception as e:
+            error_name = type(e).__name__
+            if "CredentialUnavailableError" in error_name:
+                return False, "Not authenticated - run 'az login'", None
+            return False, f"Authentication error: {e}", None
+
+
+def get_token_provider() -> AzureFoundryTokenProvider:
+    """Get the singleton token provider instance.
+
+    Returns:
+        The global AzureFoundryTokenProvider instance.
+    """
+    global _token_provider_instance
+    if _token_provider_instance is None:
+        _token_provider_instance = AzureFoundryTokenProvider()
+    return _token_provider_instance
+
+
+def reset_token_provider() -> None:
+    """Reset the singleton token provider (useful for testing)."""
+    global _token_provider_instance
+    _token_provider_instance = None