iac-code 0.1.0__py3-none-any.whl

iac_code/services/agent_factory.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import os
+import uuid
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class AgentFactoryOptions:
+    model: str
+    session_id: str | None = None
+    cwd: str | None = None
+    max_turns: int = 100
+
+
+@dataclass
+class AgentRuntime:
+    agent_loop: Any
+    session_id: str
+    tool_registry: Any
+    provider_manager: Any
+    command_registry: Any
+    task_manager: Any
+    memory_manager: Any
+
+
+def create_agent_runtime(options: AgentFactoryOptions) -> AgentRuntime:
+    from loguru import logger
+
+    from iac_code.agent.agent_loop import AgentLoop
+    from iac_code.agent.agent_tool import AgentTool
+    from iac_code.agent.system_prompt import build_system_prompt
+    from iac_code.commands import create_default_registry
+    from iac_code.commands.registry import PromptCommand
+    from iac_code.config import get_config_dir, load_credentials
+    from iac_code.memory.memory_manager import MemoryManager
+    from iac_code.memory.memory_tools import ReadMemoryTool, WriteMemoryTool
+    from iac_code.providers.manager import ProviderManager
+    from iac_code.services.cloud_credentials import CloudCredentials
+    from iac_code.services.session_storage import SessionStorage
+    from iac_code.skills.bundled import init_bundled_skills
+    from iac_code.skills.discovery import discover_all_skills, skill_to_command
+    from iac_code.skills.listing import build_skill_listing
+    from iac_code.skills.skill_tool import SkillTool
+    from iac_code.tasks.notification_queue import NotificationQueue
+    from iac_code.tasks.task_state import TaskManager
+    from iac_code.tasks.task_tools import TaskGetTool, TaskListTool, TaskStopTool
+    from iac_code.tools.base import ToolRegistry
+    from iac_code.tools.cloud.registry import register_cloud_tools
+
+    cwd = options.cwd or os.getcwd()
+    session_id = options.session_id or str(uuid.uuid4())[:8]
+
+    credentials = load_credentials()
+
+    provider_manager = ProviderManager(model=options.model, credentials=credentials)
+
+    tool_registry = ToolRegistry()
+    tool_registry.register_default_tools()
+    register_cloud_tools(tool_registry, CloudCredentials())
+
+    session_storage = SessionStorage()
+
+    memory_manager = MemoryManager(memory_dir=str(get_config_dir() / "memory"))
+    memory_content = memory_manager.get_prompt_content()
+    tool_registry.register(ReadMemoryTool(memory_manager))
+    tool_registry.register(WriteMemoryTool(memory_manager))
+
+    task_manager = TaskManager()
+    tool_registry.register(TaskListTool(task_manager))
+    tool_registry.register(TaskGetTool(task_manager))
+    tool_registry.register(TaskStopTool(task_manager))
+
+    notification_queue = NotificationQueue()
+    base_system_prompt = build_system_prompt(cwd=cwd, memory_content=memory_content)
+    tool_registry.register(
+        AgentTool(
+            task_manager=task_manager,
+            provider_manager=provider_manager,
+            tool_registry=tool_registry,
+            system_prompt=base_system_prompt,
+            notification_queue=notification_queue,
+        )
+    )
+
+    init_bundled_skills()
+    command_registry = create_default_registry()
+    for skill in discover_all_skills(cwd):
+        cmd = skill_to_command(skill)
+        existing = command_registry.get(cmd.name)
+        if existing is not None and not isinstance(existing, PromptCommand):
+            logger.warning("Skill '{}' skipped: conflicts with built-in command", cmd.name)
+            continue
+        command_registry.register(cmd)
+
+    tool_registry.register(
+        SkillTool(
+            command_registry=command_registry,
+            session_id=session_id,
+            cwd=cwd,
+            provider_manager=provider_manager,
+            tool_registry=tool_registry,
+            system_prompt=base_system_prompt,
+        )
+    )
+
+    skill_listing = build_skill_listing(command_registry.get_model_invocable_skills())
+    agent_loop = AgentLoop(
+        provider_manager=provider_manager,
+        system_prompt=build_system_prompt(cwd=cwd, memory_content=memory_content, skill_listing=skill_listing),
+        tool_registry=tool_registry,
+        session_storage=session_storage,
+        session_id=session_id,
+        max_turns=options.max_turns,
+        cwd=cwd,
+    )
+
+    return AgentRuntime(
+        agent_loop=agent_loop,
+        session_id=session_id,
+        tool_registry=tool_registry,
+        provider_manager=provider_manager,
+        command_registry=command_registry,
+        task_manager=task_manager,
+        memory_manager=memory_manager,
+    )

iac_code/services/cloud_credentials.py

@@ -0,0 +1,22 @@
+from iac_code.services.providers.aliyun import AliyunCredential, AliyunCredentials
+
+
+class CloudCredentials:
+    def __init__(self, aliyun_config_path: str | None = None) -> None:
+        self._aliyun_config_path = aliyun_config_path
+
+    def has_provider(self, name: str) -> bool:
+        if name == "aliyun":
+            return AliyunCredentials.is_configured(config_path=self._aliyun_config_path)
+        return False
+
+    def get_provider(self, name: str) -> AliyunCredential | None:
+        if name == "aliyun":
+            return AliyunCredentials.load(config_path=self._aliyun_config_path)
+        return None
+
+    def list_providers(self) -> list[str]:
+        result: list[str] = []
+        if self.has_provider("aliyun"):
+            result.append("aliyun")
+        return result

iac_code/services/context_manager.py

@@ -0,0 +1,221 @@
+"""Context manager for conversation history, token tracking, and segmented compaction."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from loguru import logger
+
+from iac_code.agent.message import ContentBlock, Conversation, Message, ToolResultBlock
+from iac_code.services.token_counter import TokenCounter
+
+
+@dataclass
+class ContextWindowConfig:
+    """Model-specific context window configuration."""
+
+    context_window: int
+    max_output_tokens: int
+    compact_buffer: int
+    compact_threshold: float
+    preserve_recent_turns: int
+
+
+_MODEL_CONFIGS: dict[str, ContextWindowConfig] = {
+    "claude": ContextWindowConfig(200_000, 8_192, 20_000, 0.93, 3),
+    "gpt-5": ContextWindowConfig(200_000, 8_192, 20_000, 0.93, 3),
+    "gpt-4": ContextWindowConfig(128_000, 8_192, 15_000, 0.93, 3),
+    "qwen": ContextWindowConfig(131_072, 8_192, 15_000, 0.93, 3),
+    "qwq": ContextWindowConfig(131_072, 8_192, 15_000, 0.93, 3),
+    "o3": ContextWindowConfig(200_000, 8_192, 20_000, 0.93, 3),
+    "o4": ContextWindowConfig(200_000, 8_192, 20_000, 0.93, 3),
+}
+_DEFAULT_CONFIG = ContextWindowConfig(128_000, 8_192, 15_000, 0.93, 3)
+
+
+def get_context_window_config(model: str) -> ContextWindowConfig:
+    model_lower = model.lower()
+    for prefix, config in _MODEL_CONFIGS.items():
+        if model_lower.startswith(prefix):
+            return config
+    return _DEFAULT_CONFIG
+
+
+class ContextManager:
+    def __init__(
+        self,
+        system_prompt: str,
+        model: str = "",
+    ) -> None:
+        self._system_prompt = system_prompt
+        self._conversation = Conversation()
+        self._model = model
+        self._token_counter = TokenCounter(model=model)
+        self._config = get_context_window_config(model)
+        self._system_prompt_tokens = self._token_counter.count_text(system_prompt)
+
+    @property
+    def system_prompt(self) -> str:
+        return self._system_prompt
+
+    @property
+    def preserve_recent_turns(self) -> int:
+        return self._config.preserve_recent_turns
+
+    @property
+    def context_window(self) -> int:
+        """Total context-window size in tokens for the current model."""
+        return self._config.context_window
+
+    def set_model(self, model: str) -> None:
+        """Switch tokenizer/context-window config for a model change.
+
+        Recomputes cached token counts so compaction thresholds stay
+        accurate after a `/model` or `/auth` switch.
+        """
+        if model == self._model:
+            return
+        self._model = model
+        self._token_counter = TokenCounter(model=model)
+        self._config = get_context_window_config(model)
+        self._system_prompt_tokens = self._token_counter.count_text(self._system_prompt)
+        for msg in self._conversation.messages:
+            msg.token_count = self._token_counter.count_message(msg.to_api_format())
+
+    def set_system_prompt(self, system_prompt: str) -> None:
+        """Replace the system prompt and refresh its cached token count."""
+        if system_prompt == self._system_prompt:
+            return
+        self._system_prompt = system_prompt
+        self._system_prompt_tokens = self._token_counter.count_text(system_prompt)
+
+    def add_user_message(self, content: str) -> Message:
+        msg = self._conversation.add_user_message(content)
+        msg.token_count = self._token_counter.count_message(msg.to_api_format())
+        return msg
+
+    def add_assistant_message(self, content: str | list[ContentBlock]) -> Message:
+        msg = self._conversation.add_assistant_message(content)
+        msg.token_count = self._token_counter.count_message(msg.to_api_format())
+        return msg
+
+    def add_tool_results(self, tool_results: list[ToolResultBlock]) -> Message:
+        msg = self._conversation.add_tool_results(tool_results)
+        msg.token_count = self._token_counter.count_message(msg.to_api_format())
+        return msg
+
+    def add_raw_message(self, raw_msg: dict[str, Any]) -> Message:
+        """Add a raw message dict (e.g. from ToolResult.new_messages) to the conversation."""
+        role = raw_msg.get("role", "user")
+        content = raw_msg.get("content", "")
+        msg = Message(role=role, content=content)
+        self._conversation.messages.append(msg)
+        msg.token_count = self._token_counter.count_message(msg.to_api_format())
+        return msg
+
+    def load_messages(self, messages: list[Message]) -> None:
+        """Inject pre-existing messages (e.g. from a resumed session)."""
+        for msg in messages:
+            self._conversation.messages.append(msg)
+            if msg.token_count == 0:
+                msg.token_count = self._token_counter.count_message(msg.to_api_format())
+
+    def get_messages(self) -> list[Message]:
+        return self._conversation.messages
+
+    def get_api_messages(self) -> list[dict[str, Any]]:
+        return self._conversation.to_api_messages()
+
+    def get_total_tokens(self) -> int:
+        return self._system_prompt_tokens + self._conversation.get_total_tokens()
+
+    def get_usage(self) -> dict[str, Any]:
+        """Return detailed token usage breakdown by category."""
+        user_tokens = 0
+        assistant_tokens = 0
+        tool_result_tokens = 0
+
+        for msg in self._conversation.messages:
+            if msg.role == "user":
+                if isinstance(msg.content, list) and any(isinstance(b, ToolResultBlock) for b in msg.content):
+                    tool_result_tokens += msg.token_count
+                else:
+                    user_tokens += msg.token_count
+            elif msg.role == "assistant":
+                assistant_tokens += msg.token_count
+
+        total = self._system_prompt_tokens + user_tokens + assistant_tokens + tool_result_tokens
+        return {
+            "system_prompt_tokens": self._system_prompt_tokens,
+            "user_message_tokens": user_tokens,
+            "assistant_message_tokens": assistant_tokens,
+            "tool_result_tokens": tool_result_tokens,
+            "total_tokens": total,
+            "context_window": self._config.context_window,
+            "usage_percent": (total / self._config.context_window * 100) if self._config.context_window > 0 else 0,
+            "message_count": len(self._conversation.messages),
+        }
+
+    def needs_compaction(self) -> bool:
+        total = self.get_total_tokens()
+        threshold = self._config.context_window * self._config.compact_threshold
+        return total > threshold
+
+    def _split_messages_for_compaction(self) -> tuple[list[Message], list[Message]]:
+        """Split messages into [old_messages, recent_messages].
+
+        A "turn" is a user+assistant message pair. We preserve the last
+        `preserve_recent_turns` turns (counting from the end).
+        """
+        messages = self._conversation.messages
+        preserve_count = self._config.preserve_recent_turns * 2
+
+        if len(messages) <= preserve_count:
+            return [], messages
+
+        split_point = len(messages) - preserve_count
+        return messages[:split_point], messages[split_point:]
+
+    def build_compaction_prompt(self) -> str:
+        """Build compaction prompt from old messages only (recent are preserved)."""
+        old_messages, _recent = self._split_messages_for_compaction()
+        if not old_messages:
+            return ""
+
+        conversation_text = []
+        for msg in old_messages:
+            role = msg.role.upper()
+            text = msg.get_text()
+            if text:
+                conversation_text.append(f"{role}: {text}")
+
+        joined = "\n".join(conversation_text)
+        return (
+            "Please provide a concise summary of this conversation so far. "
+            "Focus on:\n"
+            "1. Key decisions made\n"
+            "2. Important code changes or file modifications\n"
+            "3. Current task status and next steps\n"
+            "4. Any errors encountered and how they were resolved\n\n"
+            "Keep the summary focused and actionable. Preserve specific file paths, "
+            "function names, and technical details that are needed to continue the work.\n\n"
+            f"Conversation:\n{joined}"
+        )
+
+    def apply_compaction(self, summary: str) -> tuple[int, int]:
+        """Replace old messages with summary, keep recent messages intact."""
+        original_tokens = self._conversation.get_total_tokens()
+
+        _old, recent = self._split_messages_for_compaction()
+
+        summary_msg = Message(role="user", content=f"[Conversation Summary]\n{summary}")
+        summary_msg.token_count = self._token_counter.count_message(summary_msg.to_api_format())
+
+        self._conversation.replace_messages([summary_msg] + recent)
+        new_tokens = self._conversation.get_total_tokens()
+        logger.info(f"Compaction: {original_tokens} -> {new_tokens} tokens")
+        return (original_tokens, new_tokens)
+
+    def reset(self) -> None:
+        self._conversation = Conversation()

iac_code/services/providers/__init__.py

@@ -0,0 +1 @@
+__all__: list[str] = []

iac_code/services/providers/aliyun.py

@@ -0,0 +1,232 @@
+import json
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from iac_code.config import _load_yaml, _save_yaml, get_cloud_credentials_path
+
+DEFAULT_REGION = "cn-hangzhou"
+DEFAULT_ALIYUN_CLI_CONFIG_PATH = os.path.expanduser("~/.aliyun/config.json")
+
+# Credential modes matching aliyun CLI
+CREDENTIAL_MODES = ["AK", "StsToken", "RamRoleArn"]
+
+# Fields definition for each credential mode
+# Each field: (name, label, sensitive)
+MODE_FIELDS: dict[str, list[tuple[str, str, bool]]] = {
+    "AK": [
+        ("access_key_id", "AccessKey ID", True),
+        ("access_key_secret", "AccessKey Secret", True),
+    ],
+    "StsToken": [
+        ("access_key_id", "AccessKey ID", True),
+        ("access_key_secret", "AccessKey Secret", True),
+        ("sts_token", "STS Token", True),
+    ],
+    "RamRoleArn": [
+        ("access_key_id", "AccessKey ID", True),
+        ("access_key_secret", "AccessKey Secret", True),
+        ("ram_role_arn", "RAM Role ARN", False),
+        ("ram_session_name", "Session Name", False),
+    ],
+}
+
+# Display names for credential modes (English, translatable via i18n)
+MODE_DISPLAY_NAMES: dict[str, str] = {
+    "AK": "AccessKey",
+    "StsToken": "STS Token",
+    "RamRoleArn": "RAM Role",
+}
+
+
+@dataclass
+class AliyunCredential:
+    mode: str = "AK"
+    access_key_id: str = ""
+    access_key_secret: str = ""
+    region_id: str = field(default=DEFAULT_REGION)
+    sts_token: str = ""
+    ram_role_arn: str = ""
+    ram_session_name: str = ""
+
+
+def mask_sensitive(value: str) -> str:
+    """Mask a sensitive value with '*' characters of the same length."""
+    if not value:
+        return ""
+    return "*" * len(value)
+
+
+class AliyunCredentials:
+    @staticmethod
+    def load(config_path: str | None = None) -> AliyunCredential | None:
+        """Load credentials with priority: env vars > iac-code config > aliyun CLI config."""
+        # Try environment variables first
+        access_key_id = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_ID")
+        access_key_secret = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_SECRET")
+        if access_key_id and access_key_secret:
+            region_id = os.environ.get("ALIBABA_CLOUD_REGION_ID")
+            if not region_id:
+                # Env vars don't specify region — walk the file fallback chain:
+                # iac-code config → aliyun CLI config → DEFAULT_REGION.
+                iac_cred = AliyunCredentials._load_from_iac_code_config()
+                if iac_cred and iac_cred.region_id:
+                    region_id = iac_cred.region_id
+                else:
+                    cli_cred = AliyunCredentials._load_from_aliyun_cli(config_path)
+                    region_id = (cli_cred.region_id if cli_cred else None) or DEFAULT_REGION
+            sts_token = os.environ.get("ALIBABA_CLOUD_SECURITY_TOKEN", "")
+            mode = "StsToken" if sts_token else "AK"
+            return AliyunCredential(
+                mode=mode,
+                access_key_id=access_key_id,
+                access_key_secret=access_key_secret,
+                region_id=region_id,
+                sts_token=sts_token,
+            )
+
+        # Try iac-code config
+        if config_path is None:
+            cred = AliyunCredentials._load_from_iac_code_config()
+            if cred is not None:
+                return cred
+
+        # Fall back to aliyun CLI config
+        return AliyunCredentials._load_from_aliyun_cli(config_path)
+
+    @staticmethod
+    def _load_from_iac_code_config() -> AliyunCredential | None:
+        """Load credentials from ~/.iac-code/.cloud-credentials.yml."""
+        cloud_creds = _load_yaml(get_cloud_credentials_path())
+        aliyun_data = cloud_creds.get("aliyun")
+        if not aliyun_data or not isinstance(aliyun_data, dict):
+            return None
+
+        mode = aliyun_data.get("mode", "AK")
+        if mode not in CREDENTIAL_MODES:
+            return None
+
+        return AliyunCredential(
+            mode=mode,
+            access_key_id=aliyun_data.get("access_key_id", ""),
+            access_key_secret=aliyun_data.get("access_key_secret", ""),
+            region_id=aliyun_data.get("region_id", DEFAULT_REGION),
+            sts_token=aliyun_data.get("sts_token", ""),
+            ram_role_arn=aliyun_data.get("ram_role_arn", ""),
+            ram_session_name=aliyun_data.get("ram_session_name", ""),
+        )
+
+    @staticmethod
+    def _load_from_aliyun_cli(config_path: str | None = None) -> AliyunCredential | None:
+        """Load credentials from aliyun CLI config file (~/.aliyun/config.json)."""
+        path = Path(config_path) if config_path else Path(DEFAULT_ALIYUN_CLI_CONFIG_PATH)
+        if not path.exists():
+            return None
+
+        try:
+            data = json.loads(path.read_text())
+        except (json.JSONDecodeError, OSError):
+            return None
+
+        profiles = {p["name"]: p for p in data.get("profiles", [])}
+        profile = profiles.get("default")
+        if not profile:
+            return None
+
+        mode = profile.get("mode", "AK")
+        return AliyunCredential(
+            mode=mode,
+            access_key_id=profile.get("access_key_id", ""),
+            access_key_secret=profile.get("access_key_secret", ""),
+            region_id=profile.get("region_id", DEFAULT_REGION),
+            sts_token=profile.get("sts_token", ""),
+            ram_role_arn=profile.get("ram_role_arn", ""),
+            ram_session_name=profile.get("ram_session_name", ""),
+        )
+
+    @staticmethod
+    def load_from_aliyun_cli(config_path: str | None = None) -> AliyunCredential | None:
+        """Public method to load credentials from aliyun CLI config only.
+
+        Used by auth UI to display existing aliyun CLI config with masking.
+        """
+        return AliyunCredentials._load_from_aliyun_cli(config_path)
+
+    @staticmethod
+    def save(
+        credential: AliyunCredential,
+        config_path: str | None = None,
+    ) -> None:
+        """Save credentials to ~/.iac-code/.cloud-credentials.yml."""
+        if config_path is not None:
+            # For testing: save to specified path in aliyun CLI format
+            AliyunCredentials._save_to_aliyun_cli_format(credential, config_path)
+            return
+
+        path = get_cloud_credentials_path()
+        cloud_creds = _load_yaml(path)
+
+        aliyun_data: dict[str, Any] = {
+            "mode": credential.mode,
+            "region_id": credential.region_id,
+        }
+
+        # Save fields relevant to the credential mode
+        mode_fields = MODE_FIELDS.get(credential.mode, [])
+        for field_name, _label, _sensitive in mode_fields:
+            aliyun_data[field_name] = getattr(credential, field_name, "")
+
+        cloud_creds["aliyun"] = aliyun_data
+        _save_yaml(path, cloud_creds)
+
+    @staticmethod
+    def _save_to_aliyun_cli_format(credential: AliyunCredential, config_path: str) -> None:
+        """Save credentials in aliyun CLI JSON format (for testing)."""
+        from typing import cast
+
+        path = Path(config_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        data: dict[str, object] = {"current": "default", "profiles": []}
+        if path.exists():
+            try:
+                loaded = json.loads(path.read_text())
+                if isinstance(loaded, dict):
+                    data = loaded
+            except (json.JSONDecodeError, OSError):
+                pass
+
+        updated_profile: dict[str, str] = {
+            "name": "default",
+            "mode": credential.mode,
+            "access_key_id": credential.access_key_id,
+            "access_key_secret": credential.access_key_secret,
+            "region_id": credential.region_id,
+            "sts_token": credential.sts_token,
+            "ram_role_arn": credential.ram_role_arn,
+            "ram_session_name": credential.ram_session_name,
+        }
+
+        raw_profiles = data.get("profiles")
+        profiles: list[dict[str, str]] = (
+            cast(list[dict[str, str]], raw_profiles) if isinstance(raw_profiles, list) else []
+        )
+
+        for i, profile in enumerate(profiles):
+            if isinstance(profile, dict) and profile.get("name") == "default":
+                profiles[i] = updated_profile
+                break
+        else:
+            profiles.append(updated_profile)
+
+        data["profiles"] = profiles
+        if "current" not in data:
+            data["current"] = "default"
+
+        path.write_text(json.dumps(data, indent=2, ensure_ascii=False))
+
+    @staticmethod
+    def is_configured(config_path: str | None = None) -> bool:
+        """Check if credentials are available."""
+        return AliyunCredentials.load(config_path=config_path) is not None