airbyte-agent-mcp 0.1.33__py3-none-any.whl → 0.1.60__py3-none-any.whl

airbyte_agent_mcp/_vendored/connector_sdk/observability/config.py

@@ -0,0 +1,179 @@
+"""Unified configuration for connector-sdk."""
+
+import logging
+import os
+import tempfile
+import uuid
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+logger = logging.getLogger(__name__)
+
+# New config location
+CONFIG_DIR = Path.home() / ".airbyte" / "connector-sdk"
+CONFIG_PATH = CONFIG_DIR / "config.yaml"
+
+# Legacy file locations (for migration)
+LEGACY_USER_ID_PATH = Path.home() / ".airbyte" / "ai_sdk_user_id"
+LEGACY_INTERNAL_MARKER_PATH = Path.home() / ".airbyte" / "internal_user"
+
+
+@dataclass
+class SDKConfig:
+    """Connector SDK configuration."""
+
+    user_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    is_internal_user: bool = False
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for YAML serialization."""
+        return {
+            "user_id": self.user_id,
+            "is_internal_user": self.is_internal_user,
+        }
+
+
+def _delete_legacy_files() -> None:
+    """
+    Delete legacy config files after successful migration.
+
+    Removes:
+    - ~/.airbyte/ai_sdk_user_id
+    - ~/.airbyte/internal_user
+    """
+    for legacy_path in [LEGACY_USER_ID_PATH, LEGACY_INTERNAL_MARKER_PATH]:
+        try:
+            if legacy_path.exists():
+                legacy_path.unlink()
+                logger.debug(f"Deleted legacy config file: {legacy_path}")
+        except Exception as e:
+            logger.debug(f"Could not delete legacy file {legacy_path}: {e}")
+
+
+def _migrate_legacy_config() -> SDKConfig | None:
+    """
+    Migrate from legacy file-based config to new YAML format.
+
+    Reads from:
+    - ~/.airbyte/ai_sdk_user_id (user_id)
+    - ~/.airbyte/internal_user (is_internal_user marker)
+
+    Returns SDKConfig if migration was successful, None otherwise.
+    """
+    user_id = None
+    is_internal = False
+
+    # Try to read legacy user_id
+    try:
+        if LEGACY_USER_ID_PATH.exists():
+            user_id = LEGACY_USER_ID_PATH.read_text().strip()
+            if not user_id:
+                user_id = None
+    except Exception:
+        pass
+
+    # Check legacy internal_user marker
+    try:
+        is_internal = LEGACY_INTERNAL_MARKER_PATH.exists()
+    except Exception:
+        pass
+
+    if user_id or is_internal:
+        return SDKConfig(
+            user_id=user_id or str(uuid.uuid4()),
+            is_internal_user=is_internal,
+        )
+
+    return None
+
+
+def load_config() -> SDKConfig:
+    """
+    Load SDK configuration from config file.
+
+    Checks (in order):
+    1. New config file at ~/.airbyte/connector-sdk/config.yaml
+    2. Legacy files at ~/.airbyte/ai_sdk_user_id and ~/.airbyte/internal_user
+    3. Creates new config with generated user_id if nothing exists
+
+    Environment variable AIRBYTE_INTERNAL_USER can override is_internal_user.
+
+    Returns:
+        SDKConfig with user_id and is_internal_user
+    """
+    config = None
+
+    # Try to load from new config file
+    try:
+        if CONFIG_PATH.exists():
+            content = CONFIG_PATH.read_text()
+            data = yaml.safe_load(content) or {}
+            config = SDKConfig(
+                user_id=data.get("user_id", str(uuid.uuid4())),
+                is_internal_user=data.get("is_internal_user", False),
+            )
+            # Always clean up legacy files if they exist (even if new config exists)
+            _delete_legacy_files()
+    except Exception as e:
+        logger.debug(f"Could not load config from {CONFIG_PATH}: {e}")
+
+    # Try to migrate from legacy files if new config doesn't exist
+    if config is None:
+        config = _migrate_legacy_config()
+        if config:
+            # Save migrated config to new location
+            try:
+                save_config(config)
+                logger.debug("Migrated legacy config to new location")
+                # Delete legacy files after successful migration
+                _delete_legacy_files()
+            except Exception as e:
+                logger.debug(f"Could not save migrated config: {e}")
+
+    # Create new config if nothing exists
+    if config is None:
+        config = SDKConfig()
+        try:
+            save_config(config)
+        except Exception as e:
+            logger.debug(f"Could not save new config: {e}")
+
+    # Environment variable override for is_internal_user
+    env_value = os.getenv("AIRBYTE_INTERNAL_USER", "").lower()
+    if env_value in ("true", "1", "yes"):
+        config.is_internal_user = True
+    elif env_value:
+        # Any other non-empty value (including "false", "0", "no") defaults to False
+        config.is_internal_user = False
+
+    return config
+
+
+def save_config(config: SDKConfig) -> None:
+    """
+    Save SDK configuration to config file.
+
+    Creates the config directory if it doesn't exist.
+    Uses atomic writes to prevent corruption from concurrent access.
+
+    Args:
+        config: SDKConfig to save
+    """
+    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+
+    # Use atomic write: write to temp file then rename (atomic on POSIX)
+    fd, temp_path = tempfile.mkstemp(dir=CONFIG_DIR, suffix=".tmp")
+    try:
+        with os.fdopen(fd, "w") as f:
+            yaml.dump(config.to_dict(), f, default_flow_style=False)
+        os.rename(temp_path, CONFIG_PATH)
+    except Exception:
+        # Clean up temp file on failure
+        try:
+            os.unlink(temp_path)
+        except OSError:
+            pass
+        raise

airbyte_agent_mcp/_vendored/connector_sdk/observability/models.py

@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Dict, Optional
+from typing import Any, Dict
 
 
 @dataclass
@@ -12,8 +12,8 @@ class OperationMetadata:
     entity: str
     action: str
     timestamp: datetime
-    timing_ms: Optional[float] = None
-    status_code: Optional[int] = None
-    error_type: Optional[str] = None
-    error_message: Optional[str] = None
-    params: Optional[Dict[str, Any]] = None
+    timing_ms: float | None = None
+    status_code: int | None = None
+    error_type: str | None = None
+    error_message: str | None = None
+    params: Dict[str, Any] | None = None

airbyte_agent_mcp/_vendored/connector_sdk/observability/session.py

@@ -3,49 +3,43 @@
 import logging
 import uuid
 from datetime import UTC, datetime
-from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Dict
+
+from .config import SDKConfig, load_config
 
 logger = logging.getLogger(__name__)
 
+# Cache the config at module level to avoid repeated reads
+_cached_config: SDKConfig | None = None
+
+
+def _get_config() -> SDKConfig:
+    """Get cached SDK config or load from file."""
+    global _cached_config
+    if _cached_config is None:
+        _cached_config = load_config()
+    return _cached_config
+
+
+def _clear_config_cache() -> None:
+    """Clear the cached config. Used for testing."""
+    global _cached_config
+    _cached_config = None
+
 
 def get_persistent_user_id() -> str:
     """
-    Get or create an anonymous user ID stored in the home directory.
+    Get the persistent anonymous user ID.
 
-    The ID is stored in ~/.airbyte/ai_sdk_user_id and persists across all sessions.
-    If the file doesn't exist, a new UUID is generated and saved.
+    Now reads from ~/.airbyte/connector-sdk/config.yaml
 
     Returns:
         An anonymous UUID string that uniquely identifies this user across sessions.
     """
-    try:
-        # Create .airbyte directory in home folder if it doesn't exist
-        airbyte_dir = Path.home() / ".airbyte"
-        airbyte_dir.mkdir(exist_ok=True)
-
-        # Path to user ID file
-        user_id_file = airbyte_dir / "ai_sdk_user_id"
-
-        # Try to read existing user ID
-        if user_id_file.exists():
-            user_id = user_id_file.read_text().strip()
-            if user_id:  # Validate it's not empty
-                return user_id
+    return _get_config().user_id
 
-        # Generate new user ID if file doesn't exist or is empty
-        user_id = str(uuid.uuid4())
-        user_id_file.write_text(user_id)
-        logger.debug(f"Generated new anonymous user ID: {user_id}")
 
-        return user_id
-    except Exception as e:
-        # If we can't read/write the file, generate a session-only ID
-        logger.debug(f"Could not access anonymous user ID file: {e}")
-        return str(uuid.uuid4())
-
-
-def get_public_ip() -> Optional[str]:
+def get_public_ip() -> str | None:
     """
     Fetch the public IP address of the user.
 
@@ -53,6 +47,8 @@ def get_public_ip() -> Optional[str]:
     Uses httpx for a robust HTTP request to a public IP service.
     """
     try:
+        # NOTE: Import here intentionally - this is a non-critical network call
+        # that may fail. Importing at module level would make httpx a hard dependency.
         import httpx
 
         # Use a short timeout to avoid blocking
@@ -65,15 +61,27 @@ def get_public_ip() -> Optional[str]:
         return None
 
 
+def get_is_internal_user() -> bool:
+    """
+    Check if the current user is an internal Airbyte user.
+
+    Now reads from ~/.airbyte/connector-sdk/config.yaml
+    Environment variable AIRBYTE_INTERNAL_USER can override.
+
+    Returns False if not set or on any error.
+    """
+    return _get_config().is_internal_user
+
+
 class ObservabilitySession:
     """Shared session context for both logging and telemetry."""
 
     def __init__(
         self,
         connector_name: str,
-        connector_version: Optional[str] = None,
+        connector_version: str | None = None,
         execution_context: str = "direct",
-        session_id: Optional[str] = None,
+        session_id: str | None = None,
     ):
         self.session_id = session_id or str(uuid.uuid4())
         self.user_id = get_persistent_user_id()
@@ -84,6 +92,7 @@ class ObservabilitySession:
         self.operation_count = 0
         self.metadata: Dict[str, Any] = {}
         self.public_ip = get_public_ip()
+        self.is_internal_user = get_is_internal_user()
 
     def increment_operations(self):
         """Increment the operation counter."""

airbyte_agent_mcp/_vendored/connector_sdk/performance/metrics.py

@@ -2,7 +2,7 @@
 
 import time
 from contextlib import asynccontextmanager
-from typing import Dict, Optional
+from typing import Dict
 
 
 class PerformanceMonitor:
@@ -33,7 +33,7 @@ class PerformanceMonitor:
         metrics["min"] = min(metrics["min"], duration)
         metrics["max"] = max(metrics["max"], duration)
 
-    def get_stats(self, metric_name: str) -> Optional[Dict[str, float]]:
+    def get_stats(self, metric_name: str) -> Dict[str, float] | None:
         """Get statistics for a metric.
 
         Args:
@@ -62,7 +62,7 @@ class PerformanceMonitor:
         """
         return {name: self.get_stats(name) for name in self._metrics.keys()}
 
-    def reset(self, metric_name: Optional[str] = None):
+    def reset(self, metric_name: str | None = None):
         """Reset metrics.
 
         Args:

airbyte_agent_mcp/_vendored/connector_sdk/schema/base.py

@@ -7,13 +7,13 @@ References:
 """
 
 from enum import StrEnum
-from typing import Dict, Optional
+from typing import Dict
 from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 from pydantic_core import Url
 
-from .extensions import RetryConfig
+from .extensions import CacheConfig, RetryConfig
 
 
 class ExampleQuestions(BaseModel):
@@ -45,9 +45,9 @@ class Contact(BaseModel):
 
     model_config = ConfigDict(populate_by_name=True, extra="forbid")
 
-    name: Optional[str] = None
-    url: Optional[str] = None
-    email: Optional[str] = None
+    name: str | None = None
+    url: str | None = None
+    email: str | None = None
 
 
 class License(BaseModel):
@@ -60,7 +60,7 @@ class License(BaseModel):
     model_config = ConfigDict(populate_by_name=True, extra="forbid")
 
     name: str
-    url: Optional[str] = None
+    url: str | None = None
 
 
 class DocUrlType(StrEnum):
@@ -85,7 +85,7 @@ class DocUrl(BaseModel):
 
     url: str
     type: DocUrlType
-    title: Optional[str] = None
+    title: str | None = None
 
     @field_validator("url")
     def validate_url(cls, v):
@@ -105,23 +105,25 @@ class Info(BaseModel):
     - x-airbyte-external-documentation-urls: List of external documentation URLs (Airbyte extension)
     - x-airbyte-retry-config: Retry configuration for transient errors (Airbyte extension)
     - x-airbyte-example-questions: Example questions for AI connector README (Airbyte extension)
+    - x-airbyte-cache: Cache configuration for field mapping between API and cache schemas (Airbyte extension)
     """
 
     model_config = ConfigDict(populate_by_name=True, extra="forbid")
 
     title: str
     version: str
-    description: Optional[str] = None
-    terms_of_service: Optional[str] = Field(None, alias="termsOfService")
-    contact: Optional[Contact] = None
-    license: Optional[License] = None
+    description: str | None = None
+    terms_of_service: str | None = Field(None, alias="termsOfService")
+    contact: Contact | None = None
+    license: License | None = None
 
     # Airbyte extension
-    x_airbyte_connector_name: Optional[str] = Field(None, alias="x-airbyte-connector-name")
-    x_airbyte_connector_id: Optional[UUID] = Field(None, alias="x-airbyte-connector-id")
+    x_airbyte_connector_name: str | None = Field(None, alias="x-airbyte-connector-name")
+    x_airbyte_connector_id: UUID | None = Field(None, alias="x-airbyte-connector-id")
     x_airbyte_external_documentation_urls: list[DocUrl] = Field(..., alias="x-airbyte-external-documentation-urls")
-    x_airbyte_retry_config: Optional[RetryConfig] = Field(None, alias="x-airbyte-retry-config")
-    x_airbyte_example_questions: Optional[ExampleQuestions] = Field(None, alias="x-airbyte-example-questions")
+    x_airbyte_retry_config: RetryConfig | None = Field(None, alias="x-airbyte-retry-config")
+    x_airbyte_example_questions: ExampleQuestions | None = Field(None, alias="x-airbyte-example-questions")
+    x_airbyte_cache: CacheConfig | None = Field(None, alias="x-airbyte-cache")
 
 
 class ServerVariable(BaseModel):
@@ -133,9 +135,9 @@ class ServerVariable(BaseModel):
 
     model_config = ConfigDict(populate_by_name=True, extra="forbid")
 
-    enum: Optional[list[str]] = None
+    enum: list[str] | None = None
     default: str
-    description: Optional[str] = None
+    description: str | None = None
 
 
 class Server(BaseModel):
@@ -148,7 +150,7 @@ class Server(BaseModel):
     model_config = ConfigDict(populate_by_name=True, extra="forbid")
 
     url: str
-    description: Optional[str] = None
+    description: str | None = None
     variables: Dict[str, ServerVariable] = Field(default_factory=dict)
 
     @field_validator("url")

airbyte_agent_mcp/_vendored/connector_sdk/schema/components.py

@@ -7,7 +7,7 @@ References:
 - https://spec.openapis.org/oas/v3.1.0#parameter-object
 """
 
-from typing import Any, Dict, List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Union
 
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -30,43 +30,44 @@ class Schema(BaseModel):
     model_config = ConfigDict(populate_by_name=True, extra="forbid")
 
     # Core JSON Schema fields
-    type: Optional[str] = None
-    format: Optional[str] = None
-    title: Optional[str] = None
-    description: Optional[str] = None
-    default: Optional[Any] = None
-    example: Optional[Any] = None
+    type: str | None = None
+    format: str | None = None
+    title: str | None = None
+    description: str | None = None
+    default: Any | None = None
+    example: Any | None = None
 
     # Object properties
     properties: Dict[str, Any] = Field(default_factory=dict)  # May contain $ref
     required: List[str] = Field(default_factory=list)
-    additional_properties: Optional[Any] = Field(None, alias="additionalProperties")
+    additional_properties: Any | None = Field(None, alias="additionalProperties")
 
     # Array properties
-    items: Optional[Any] = None  # May be Schema or $ref
+    items: Any | None = None  # May be Schema or $ref
 
     # Validation
-    enum: Optional[List[Any]] = None
-    min_length: Optional[int] = Field(None, alias="minLength")
-    max_length: Optional[int] = Field(None, alias="maxLength")
-    minimum: Optional[float] = None
-    maximum: Optional[float] = None
-    pattern: Optional[str] = None
+    enum: List[Any] | None = None
+    min_length: int | None = Field(None, alias="minLength")
+    max_length: int | None = Field(None, alias="maxLength")
+    minimum: float | None = None
+    maximum: float | None = None
+    pattern: str | None = None
 
     # Composition
-    all_of: Optional[List[Any]] = Field(None, alias="allOf")
-    any_of: Optional[List[Any]] = Field(None, alias="anyOf")
-    one_of: Optional[List[Any]] = Field(None, alias="oneOf")
-    not_: Optional[Any] = Field(None, alias="not")
+    all_of: List[Any] | None = Field(None, alias="allOf")
+    any_of: List[Any] | None = Field(None, alias="anyOf")
+    one_of: List[Any] | None = Field(None, alias="oneOf")
+    not_: Any | None = Field(None, alias="not")
 
     # Metadata
-    nullable: Optional[bool] = Field(None, deprecated="Use type union with null instead (OpenAPI 3.1)")
-    read_only: Optional[bool] = Field(None, alias="readOnly")
-    write_only: Optional[bool] = Field(None, alias="writeOnly")
-    deprecated: Optional[bool] = None
+    nullable: bool | None = Field(None, deprecated="Use type union with null instead (OpenAPI 3.1)")
+    read_only: bool | None = Field(None, alias="readOnly")
+    write_only: bool | None = Field(None, alias="writeOnly")
+    deprecated: bool | None = None
 
-    # Airbyte extension
-    x_airbyte_entity_name: Optional[str] = Field(None, alias="x-airbyte-entity-name")
+    # Airbyte extensions
+    x_airbyte_entity_name: str | None = Field(None, alias="x-airbyte-entity-name")
+    x_airbyte_stream_name: str | None = Field(None, alias="x-airbyte-stream-name")
 
 
 class Parameter(BaseModel):
@@ -80,19 +81,19 @@ class Parameter(BaseModel):
 
     name: str
     in_: Literal["query", "header", "path", "cookie"] = Field(alias="in")
-    description: Optional[str] = None
-    required: Optional[bool] = None
-    deprecated: Optional[bool] = None
-    allow_empty_value: Optional[bool] = Field(None, alias="allowEmptyValue")
+    description: str | None = None
+    required: bool | None = None
+    deprecated: bool | None = None
+    allow_empty_value: bool | None = Field(None, alias="allowEmptyValue")
 
     # Schema can be inline or reference
-    schema_: Optional[Dict[str, Any]] = Field(None, alias="schema")
+    schema_: Dict[str, Any] | None = Field(None, alias="schema")
 
     # Style and examples
-    style: Optional[str] = None
-    explode: Optional[bool] = None
-    example: Optional[Any] = None
-    examples: Optional[Dict[str, Any]] = None
+    style: str | None = None
+    explode: bool | None = None
+    example: Any | None = None
+    examples: Dict[str, Any] | None = None
 
 
 class MediaType(BaseModel):
@@ -104,10 +105,10 @@ class MediaType(BaseModel):
 
     model_config = ConfigDict(populate_by_name=True, extra="forbid")
 
-    schema_: Optional[Dict[str, Any]] = Field(None, alias="schema")
-    example: Optional[Any] = None
-    examples: Optional[Dict[str, Any]] = None
-    encoding: Optional[Dict[str, Any]] = None
+    schema_: Dict[str, Any] | None = Field(None, alias="schema")
+    example: Any | None = None
+    examples: Dict[str, Any] | None = None
+    encoding: Dict[str, Any] | None = None
 
 
 class GraphQLBodyConfig(BaseModel):
@@ -124,12 +125,12 @@ class GraphQLBodyConfig(BaseModel):
         ...,
         description="GraphQL query or mutation string with optional template placeholders (e.g., {{ variable }})",
     )
-    variables: Optional[Dict[str, Any]] = Field(
+    variables: Dict[str, Any] | None = Field(
         None,
         description="Variables to substitute in the GraphQL query using template syntax (e.g., {{ param_name }})",
     )
-    operationName: Optional[str] = Field(None, description="Operation name for queries with multiple operations")
-    default_fields: Optional[Union[str, List[str]]] = Field(
+    operationName: str | None = Field(None, description="Operation name for queries with multiple operations")
+    default_fields: Union[str, List[str]] | None = Field(
         None,
         description="Default fields to select if not provided in request parameters. Can be a string or array of field names.",
     )
@@ -155,7 +156,7 @@ class PathOverrideConfig(BaseModel):
 
     path: str = Field(
         ...,
-        description=("Actual HTTP path to use for requests (e.g., '/graphql'). " "Must start with '/'"),
+        description=("Actual HTTP path to use for requests (e.g., '/graphql'). Must start with '/'"),
     )
 
 
@@ -172,17 +173,17 @@ class RequestBody(BaseModel):
 
     model_config = ConfigDict(populate_by_name=True, extra="forbid")
 
-    description: Optional[str] = None
+    description: str | None = None
     content: Dict[str, MediaType] = Field(default_factory=dict)
-    required: Optional[bool] = None
+    required: bool | None = None
 
     # Airbyte extensions for GraphQL support
     # See connector_sdk.extensions for AIRBYTE_BODY_TYPE constant
-    x_airbyte_body_type: Optional[BodyTypeConfig] = Field(
+    x_airbyte_body_type: BodyTypeConfig | None = Field(
         None,
         alias="x-airbyte-body-type",  # AIRBYTE_BODY_TYPE
         description=(
-            "Body type and configuration. Contains 'type' field (e.g., 'graphql') " "and type-specific configuration (query, variables, etc.)."
+            "Body type and configuration. Contains 'type' field (e.g., 'graphql') and type-specific configuration (query, variables, etc.)."
         ),
     )
 
@@ -196,11 +197,11 @@ class Header(BaseModel):
 
     model_config = ConfigDict(populate_by_name=True, extra="forbid")
 
-    description: Optional[str] = None
-    required: Optional[bool] = None
-    deprecated: Optional[bool] = None
-    schema_: Optional[Dict[str, Any]] = Field(None, alias="schema")
-    example: Optional[Any] = None
+    description: str | None = None
+    required: bool | None = None
+    deprecated: bool | None = None
+    schema_: Dict[str, Any] | None = Field(None, alias="schema")
+    example: Any | None = None
 
 
 class Response(BaseModel):
@@ -213,9 +214,9 @@ class Response(BaseModel):
     model_config = ConfigDict(populate_by_name=True, extra="forbid")
 
     description: str
-    headers: Optional[Dict[str, Header]] = None
-    content: Optional[Dict[str, MediaType]] = None
-    links: Optional[Dict[str, Any]] = None
+    headers: Dict[str, Header] | None = None
+    content: Dict[str, MediaType] | None = None
+    links: Dict[str, Any] | None = None
 
 
 class Components(BaseModel):
@@ -230,9 +231,9 @@ class Components(BaseModel):
     schemas: Dict[str, Schema] = Field(default_factory=dict)
     responses: Dict[str, Response] = Field(default_factory=dict)
     parameters: Dict[str, Parameter] = Field(default_factory=dict)
-    examples: Optional[Dict[str, Any]] = None
+    examples: Dict[str, Any] | None = None
     request_bodies: Dict[str, RequestBody] = Field(default_factory=dict, alias="requestBodies")
-    headers: Optional[Dict[str, Header]] = None
+    headers: Dict[str, Header] | None = None
     security_schemes: Dict[str, SecurityScheme] = Field(default_factory=dict, alias="securitySchemes")
-    links: Optional[Dict[str, Any]] = None
-    callbacks: Optional[Dict[str, Any]] = None
+    links: Dict[str, Any] | None = None
+    callbacks: Dict[str, Any] | None = None