augment-sdk 0.1.1__py3-none-any.whl

augment/agent.py

@@ -0,0 +1,1139 @@
+"""
+Main Agent class for the Augment Python SDK.
+
+This implementation uses the ACP (Agent Client Protocol) client for
+better performance and real-time streaming.
+"""
+
+import json
+import os
+import re
+import subprocess
+from contextlib import contextmanager
+from dataclasses import dataclass, fields, is_dataclass
+from enum import Enum
+from pathlib import Path
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    get_args,
+    get_origin,
+)
+
+from .acp import ACPClient, AuggieACPClient
+from .exceptions import (
+    AugmentCLIError,
+    AugmentNotFoundError,
+    AugmentParseError,
+    AugmentWorkspaceError,
+    AugmentVerificationError,
+)
+from .listener import AgentListener
+from .listener_adapter import AgentListenerAdapter
+from .prompt_formatter import AgentPromptFormatter, DEFAULT_INFERENCE_TYPES
+
+T = TypeVar("T")
+
+
+@dataclass
+class VerificationResult:
+    """Result of verifying success criteria."""
+
+    all_criteria_met: bool
+    unmet_criteria: List[int]  # Indices of criteria that are not met (1-based)
+    issues: List[str]  # Description of each issue found
+    overall_assessment: str  # Brief overall assessment
+
+
+def _get_type_name(type_hint: Any) -> str:
+    """
+    Safely get the name of a type hint.
+
+    Handles both actual types and string annotations.
+
+    Args:
+        type_hint: A type or type annotation
+
+    Returns:
+        String name of the type
+    """
+    if isinstance(type_hint, str):
+        return type_hint
+    if hasattr(type_hint, "__name__"):
+        return type_hint.__name__
+    return str(type_hint)
+
+
+@dataclass
+class Model:
+    """
+    Represents an available AI model.
+
+    Attributes:
+        id: The model identifier used with --model flag (e.g., "sonnet4.5")
+        name: The human-readable model name (e.g., "Claude Sonnet 4.5")
+        description: Additional information about the model (e.g., "Anthropic Claude Sonnet 4.5, 200k context")
+    """
+
+    id: str
+    name: str
+    description: str
+
+    def __str__(self) -> str:
+        return f"{self.name} [{self.id}]"
+
+
+class Agent:
+    """
+    Augment CLI agent interface for programmatic access.
+
+    This class provides a Python interface to the Augment CLI agent (auggie),
+    using the ACP (Agent Client Protocol) for better performance and
+    real-time streaming of responses.
+
+    By default, each run() call creates a fresh session. Use the session()
+    context manager to maintain conversation continuity across multiple calls.
+
+    Attributes:
+        last_model_answer: The last textual explanation returned by the model
+            when using typed results. This contains the agent's reasoning or
+            context and may be helpful for logging and debugging. None for
+            untyped responses or if no message was provided.
+        model: The AI model to use for instructions. None uses the CLI's default.
+        workspace_path: The resolved workspace path for this agent.
+        session_id: The current session ID (only set when inside a session context).
+        timeout: Default timeout in seconds for agent operations (defaults to 180 seconds).
+    """
+
+    def __init__(
+        self,
+        workspace_root: Optional[Union[str, Path]] = None,
+        model: Optional[str] = None,
+        listener: Optional[AgentListener] = None,
+        cli_path: Optional[str] = None,
+        acp_client: Optional[ACPClient] = None,
+        removed_tools: Optional[List[str]] = None,
+        timeout: int = 180,
+        api_key: Optional[str] = None,
+        api_url: Optional[str] = None,
+    ):
+        """
+        Initialize an agent instance.
+
+        Args:
+            workspace_root: Path to the workspace root. Defaults to current directory.
+            model: The AI model to use (e.g., "claude-3-5-sonnet-latest", "gpt-4o").
+                   Defaults to the CLI's default model.
+            listener: Optional listener for agent events (AgentListener).
+            cli_path: Optional path to the Augment CLI. Auto-detected if not provided.
+            acp_client: Optional ACP client instance for testing. If provided, this client
+                       will be used instead of creating a new one. This is primarily for
+                       testing purposes to allow mocking the ACP client.
+            removed_tools: List of tool names to remove/disable (e.g., ["github-api", "linear"]).
+                          These tools will not be available to the agent.
+            timeout: Default timeout in seconds for agent operations. Defaults to 180 seconds (3 minutes).
+                    This timeout is used when no timeout is specified in the run() method.
+            api_key: Optional API key for authentication. If provided, sets AUGMENT_API_TOKEN
+                    environment variable for the agent process.
+            api_url: Optional API URL. If not provided, uses AUGMENT_API_URL environment variable,
+                    or defaults to "https://api.augmentcode.com". Sets AUGMENT_API_URL environment
+                    variable for the agent process.
+        """
+        self.workspace_path = self._validate_workspace_path(workspace_root)
+        self.model = model
+        self.listener = listener
+        # Create adapter to bridge ACP events to AgentListener
+        self._listener_adapter = AgentListenerAdapter(listener) if listener else None
+        self.cli_path = cli_path
+        self.removed_tools = removed_tools or []
+        self.timeout = timeout
+        self.api_key = api_key
+        self.api_url = (
+            api_url
+            if api_url is not None
+            else os.getenv("AUGMENT_API_URL", "https://api.augmentcode.com")
+        )
+        self.last_model_answer: Optional[str] = None
+        self._acp_client: Optional[ACPClient] = (
+            acp_client  # Use provided client or None
+        )
+        self._provided_client = acp_client is not None  # Track if client was provided
+        self._in_session = False  # Track if we're in a session context
+        self._prompt_formatter = AgentPromptFormatter()  # Handles prompt formatting
+
+    def run(
+        self,
+        instruction: str,
+        return_type: Optional[Type[T]] = None,
+        timeout: Optional[int] = None,
+        max_retries: int = 3,
+        functions: Optional[List[Callable]] = None,
+        success_criteria: Optional[List[str]] = None,
+        max_verification_rounds: int = 3,
+    ) -> Union[T, Any]:
+        """
+        Execute an instruction and return the agent's response.
+
+        If a typed response is requested and parsing fails, the agent will be asked
+        to correct its output up to max_retries times.
+
+        When no return_type is specified, the agent automatically infers the best type
+        from common types (int, float, bool, str, list, dict) and returns the result.
+
+        Args:
+            instruction: The instruction to send to the agent
+            return_type: Expected return type (int, str, bool, float, list, dict, dataclass, or enum).
+                        If None, the agent will automatically infer the type.
+            timeout: Optional timeout in seconds
+            max_retries: Maximum number of retry attempts for parsing failures (default: 3)
+            functions: Optional list of Python functions the agent can call. Functions should
+                      have type hints and docstrings describing their parameters.
+            success_criteria: Optional list of criteria that must be met after execution.
+                            The agent will iteratively work on the task and verify criteria
+                            until all are met or max_verification_rounds is reached.
+            max_verification_rounds: Maximum number of verification rounds when using
+                                   success_criteria (default: 3). Each round consists of
+                                   executing/fixing the task and verifying criteria.
+
+        Returns:
+            Parsed result (either of requested type, or inferred type)
+
+        Note:
+            After execution, self.last_model_answer contains the agent's textual explanation,
+            which may be helpful for logging and debugging. This is separate from the
+            structured return value and provides the model's reasoning or context.
+
+        Raises:
+            ValueError: If instruction is empty or whitespace, or unsupported return_type
+            AugmentParseError: If parsing typed response fails after all retries
+            AugmentVerificationError: If success_criteria are not met after max_verification_rounds
+        """
+        # Validate instruction
+        if not instruction or not instruction.strip():
+            raise ValueError("Instruction cannot be empty or whitespace")
+
+        # Prepare the client
+        client = self._prepare_client()
+
+        # If no success criteria, just run the task once
+        if not success_criteria:
+            return self._run_task(
+                client=client,
+                instruction=instruction,
+                return_type=return_type,
+                timeout=timeout,
+                max_retries=max_retries,
+                functions=functions,
+            )
+
+        # With success criteria: loop until verified or max rounds
+        current_instruction = instruction
+
+        for round_num in range(max_verification_rounds):
+            # Execute the task (or fix issues from previous round)
+            result = self._run_task(
+                client=client,
+                instruction=current_instruction,
+                return_type=return_type,
+                timeout=timeout,
+                max_retries=max_retries,
+                functions=functions,
+            )
+
+            # Verify success criteria
+            verification = self._verify_success_criteria(
+                client=client,
+                success_criteria=success_criteria,
+                timeout=timeout,
+            )
+
+            # Check if all criteria are met
+            if verification.all_criteria_met:
+                # Success! All criteria verified
+                return result
+
+            # Not all criteria met - prepare fix instruction for next round
+            if round_num < max_verification_rounds - 1:
+                # We have more rounds - prepare fix instruction
+                current_instruction = self._prepare_fix_instruction(
+                    success_criteria=success_criteria,
+                    verification=verification,
+                )
+            else:
+                # Last round - raise exception
+                raise AugmentVerificationError(
+                    f"Success criteria not fully met after {max_verification_rounds} rounds. "
+                    f"Unmet criteria: {verification.unmet_criteria}. "
+                    f"Issues: {verification.issues}",
+                    unmet_criteria=verification.unmet_criteria,
+                    issues=verification.issues,
+                    rounds_attempted=max_verification_rounds,
+                )
+
+        # This should never be reached, but just in case
+        return None
+
+    def _prepare_client(self) -> ACPClient:
+        """
+        Prepare the ACP client for execution.
+
+        This method:
+        1. Creates the client if it doesn't exist
+        2. Starts the client if it's not running
+        3. Clears context if not in a session (for fresh start)
+
+        Returns:
+            ACP client instance ready for use
+        """
+        # Create client if needed
+        if self._acp_client is None:
+            self._acp_client = AuggieACPClient(
+                cli_path=self.cli_path,
+                model=self.model,
+                workspace_root=str(self.workspace_path),
+                listener=self._listener_adapter,  # Pass the adapter to ACP
+                removed_tools=self.removed_tools,
+                api_key=self.api_key,
+                api_url=self.api_url,
+            )
+
+        # Start if not running
+        if not self._acp_client.is_running:
+            self._acp_client.start()
+
+        # If not in a session, clear context for a fresh start
+        if not self._in_session:
+            self._acp_client.clear_context()
+
+        return self._acp_client
+
+    def _run_task(
+        self,
+        client: ACPClient,
+        instruction: str,
+        return_type: Optional[Type[T]],
+        timeout: Optional[int],
+        max_retries: int,
+        functions: Optional[List[Callable]] = None,
+    ) -> Union[T, Any]:
+        """
+        Run a single task: talk to the agent in a loop until the task is complete.
+
+        This method implements a loop that:
+        1. Sends the current instruction to the agent
+        2. If response contains function calls, executes them and continues
+        3. If response can be parsed successfully, returns the result
+        4. If parsing fails and retries remain, asks agent to fix output and continues
+        5. Repeats until success or max rounds exhausted
+
+        Args:
+            client: The ACP client to use
+            instruction: The instruction to send
+            return_type: Expected return type (None for type inference)
+            timeout: Timeout in seconds (None for default)
+            max_retries: Maximum number of retry attempts for parsing failures
+            functions: Optional list of functions the agent can call
+
+        Returns:
+            Parsed result (either of requested type, or inferred type)
+
+        Raises:
+            AugmentParseError: If parsing fails after all retries
+        """
+        # Set effective timeout - use provided timeout, or fall back to instance default
+        effective_timeout = (
+            float(timeout) if timeout is not None else float(self.timeout)
+        )
+
+        # Prepare the instruction prompt
+        current_instruction, function_map = (
+            self._prompt_formatter.prepare_instruction_prompt(
+                instruction, return_type, functions
+            )
+        )
+
+        # Determine if we're in type inference mode
+        infer_type = return_type is None
+
+        # Track retry attempts for parsing failures
+        parse_retry_count = 0
+
+        # Maximum total rounds to prevent infinite loops
+        # This includes both function calls and parse retries
+        max_rounds = 20
+
+        for round_num in range(max_rounds):
+            # Send message to agent
+            raw_response = client.send_message(current_instruction, effective_timeout)
+
+            # Check if there are function calls to handle
+            if function_map:
+                function_calls = self._parse_function_calls(raw_response)
+                if function_calls:
+                    # Execute functions and prepare next instruction with results
+                    current_instruction = self._handle_function_calls(
+                        raw_response, function_map
+                    )
+                    # Continue loop to send function results back to agent
+                    continue
+
+            # No function calls - try to parse the response
+            try:
+                if infer_type:
+                    result, _ = self._parse_type_inference_response(
+                        raw_response, DEFAULT_INFERENCE_TYPES
+                    )
+                    return result
+                else:
+                    # At this point, return_type must be a valid type (not None)
+                    assert return_type is not None
+                    result = self._parse_typed_response(raw_response, return_type)
+                    return result  # type: ignore[return-value,no-any-return]
+
+            except AugmentParseError as e:
+                # Parsing failed - check if we have retries left
+                if parse_retry_count >= max_retries:
+                    raise
+
+                # Increment retry count and prepare retry instruction
+                parse_retry_count += 1
+                current_instruction = self._prompt_formatter.prepare_retry_prompt(
+                    return_type,
+                    str(e),
+                    parse_retry_count,
+                )
+                # Continue loop to retry with corrected instruction
+                continue
+
+        # If we've exhausted all rounds, raise an error
+        raise AugmentParseError(
+            f"Exceeded maximum rounds ({max_rounds}) without successful completion"
+        )
+
+    def _verify_success_criteria(
+        self,
+        client: ACPClient,
+        success_criteria: List[str],
+        timeout: Optional[int],
+    ) -> VerificationResult:
+        """
+        Verify that success criteria are met and return structured feedback.
+
+        This method checks each criterion and returns detailed information about
+        which criteria are met and which are not, along with specific issues.
+
+        Args:
+            client: The ACP client to use
+            success_criteria: List of criteria that must be met
+            timeout: Timeout in seconds (None for default)
+
+        Returns:
+            VerificationResult with detailed verification feedback
+        """
+        # Build verification instruction
+        criteria_list = "\n".join(
+            f"{i + 1}. {criterion}" for i, criterion in enumerate(success_criteria)
+        )
+
+        verification_instruction = f"""The task has been completed. Please verify that ALL of the following success criteria are met:
+
+{criteria_list}
+
+For each criterion, check if it is satisfied.
+
+Respond with a JSON object with the following structure:
+{{
+    "all_criteria_met": true/false,
+    "unmet_criteria": [list of criterion numbers (1-based) that are NOT met],
+    "issues": ["description of each issue found"],
+    "overall_assessment": "brief assessment of the current state"
+}}
+
+Example:
+{{
+    "all_criteria_met": false,
+    "unmet_criteria": [2, 3],
+    "issues": ["Criterion 2: Function is missing type hints for parameter 'x'", "Criterion 3: No docstring present"],
+    "overall_assessment": "Function exists but lacks type hints and documentation"
+}}"""
+
+        # Run verification and get structured result
+        try:
+            verification = self._run_task(
+                client=client,
+                instruction=verification_instruction,
+                return_type=VerificationResult,
+                timeout=timeout,
+                max_retries=2,
+                functions=None,
+            )
+            return verification
+
+        except (AugmentParseError, Exception) as e:
+            # If verification fails, return a conservative result
+            import warnings
+
+            warnings.warn(
+                f"Success criteria verification failed: {e}. "
+                "Assuming criteria are not met.",
+                UserWarning,
+            )
+            return VerificationResult(
+                all_criteria_met=False,
+                unmet_criteria=list(range(1, len(success_criteria) + 1)),
+                issues=[f"Verification failed with error: {e}"],
+                overall_assessment="Verification could not be completed",
+            )
+
+    def _prepare_fix_instruction(
+        self,
+        success_criteria: List[str],
+        verification: VerificationResult,
+    ) -> str:
+        """
+        Prepare an instruction to fix issues identified during verification.
+
+        Args:
+            success_criteria: List of all success criteria
+            verification: Verification result with issues
+
+        Returns:
+            Instruction for the agent to fix the issues
+        """
+        # Build list of unmet criteria with their descriptions
+        unmet_details = []
+        for criterion_num in verification.unmet_criteria:
+            if 1 <= criterion_num <= len(success_criteria):
+                criterion_text = success_criteria[criterion_num - 1]
+                unmet_details.append(f"{criterion_num}. {criterion_text}")
+
+        unmet_list = "\n".join(unmet_details)
+        issues_list = "\n".join(f"- {issue}" for issue in verification.issues)
+
+        fix_instruction = f"""The following success criteria are NOT yet met:
+
+{unmet_list}
+
+Issues identified:
+{issues_list}
+
+Overall assessment: {verification.overall_assessment}
+
+Please fix these issues to ensure ALL success criteria are satisfied."""
+
+        return fix_instruction
+
+    def _handle_function_calls(
+        self,
+        raw_response: str,
+        function_map: Dict[str, Callable],
+    ) -> str:
+        """
+        Execute function calls from the agent's response and prepare next instruction.
+
+        This method:
+        1. Parses function calls from the response
+        2. Executes the functions
+        3. Returns an instruction containing the function results to send back to the agent
+
+        Args:
+            raw_response: The agent's response that contains function calls
+            function_map: Mapping of function names to callables
+
+        Returns:
+            Instruction text with function results to send back to the agent
+        """
+        # Parse function calls from response
+        function_calls = self._parse_function_calls(raw_response)
+
+        # Execute function calls
+        function_results = []
+        for func_call in function_calls:
+            func_name = func_call.get("name")
+            func_args = func_call.get("arguments", {})
+
+            # Skip if function name is missing
+            if not func_name:
+                continue
+
+            if func_name not in function_map:
+                result = {"error": f"Function '{func_name}' not found"}
+                error = f"Function '{func_name}' not found"
+                # Notify listener of error
+                if self.listener:
+                    self.listener.on_function_result(func_name, None, error)
+            else:
+                try:
+                    func = function_map[func_name]
+
+                    # Notify listener of function call (right before execution)
+                    if self.listener:
+                        self.listener.on_function_call(func_name, func_args)
+
+                    result = func(**func_args)
+                    # Notify listener of success
+                    if self.listener:
+                        self.listener.on_function_result(func_name, result, None)
+                except Exception as e:
+                    result = {"error": f"Error calling {func_name}: {str(e)}"}
+                    # Notify listener of error
+                    if self.listener:
+                        self.listener.on_function_result(func_name, None, str(e))
+
+            function_results.append({"function": func_name, "result": result})
+
+        # Build follow-up instruction with function results
+        results_text = "Function call results:\n\n"
+        for fr in function_results:
+            results_text += f"Function: {fr['function']}\n"
+            results_text += f"Result: {json.dumps(fr['result'], indent=2)}\n\n"
+
+        results_text += "\nPlease continue with your response based on these results."
+
+        return results_text
+
+    def _parse_function_calls(self, response: str) -> List[Dict[str, Any]]:
+        """
+        Parse function calls from agent response.
+
+        Looks for <function-call>...</function-call> blocks containing JSON.
+
+        Args:
+            response: Agent's response text
+
+        Returns:
+            List of function call dictionaries with 'name' and 'arguments' keys
+        """
+        function_calls = []
+        pattern = r"<function-call>\s*(\{.*?\})\s*</function-call>"
+        matches = re.findall(pattern, response, re.DOTALL)
+
+        for match in matches:
+            try:
+                func_call = json.loads(match)
+                if "name" in func_call:
+                    function_calls.append(func_call)
+            except json.JSONDecodeError:
+                # Skip invalid JSON
+                continue
+
+        return function_calls
+
+    def _parse_type_inference_response(
+        self, response: str, possible_types: List[Type[Any]]
+    ) -> Tuple[Any, Type[Any]]:
+        """
+        Parse type inference response.
+
+        Args:
+            response: The agent's response
+            possible_types: List of possible types
+
+        Returns:
+            Tuple of (parsed result, chosen type)
+
+        Raises:
+            AugmentParseError: If parsing fails
+        """
+        # Extract message
+        message_match = re.search(
+            r"<augment-agent-message>\s*(.*?)\s*</augment-agent-message>",
+            response,
+            re.DOTALL,
+        )
+        if message_match:
+            self.last_model_answer = message_match.group(1).strip()
+        else:
+            self.last_model_answer = None
+
+        # Extract type name
+        type_match = re.search(
+            r"<augment-agent-type>\s*(\w+)\s*</augment-agent-type>", response, re.DOTALL
+        )
+        if not type_match:
+            # If no type tags found, check if response is empty or just whitespace
+            # This can happen when agent completes a task (like file creation) successfully
+            # but doesn't provide a structured response
+            if not response or response.strip() == "":
+                # Return empty string as success indicator
+                return "", str
+
+            # If there's content but no tags, try to extract it as a string
+            # This handles cases where agent responds without proper formatting
+            content = response.strip()
+            if content:
+                # Return the content as a string
+                return content, str
+
+            raise AugmentParseError(
+                "No type classification found. Expected <augment-agent-type> tags."
+            )
+
+        type_name = type_match.group(1).strip()
+
+        # Find the matching type by name
+        chosen_type = None
+        for t in possible_types:
+            if t.__name__ == type_name:
+                chosen_type = t
+                break
+
+        if chosen_type is None:
+            type_names = [t.__name__ for t in possible_types]
+            raise AugmentParseError(
+                f"Invalid type name '{type_name}'. Must be one of: {', '.join(type_names)}"
+            )
+
+        # Extract result
+        result_match = re.search(
+            r"<augment-agent-result>\s*(.*?)\s*</augment-agent-result>",
+            response,
+            re.DOTALL,
+        )
+        if not result_match:
+            raise AugmentParseError(
+                "No structured result found. Expected <augment-agent-result> tags."
+            )
+
+        content = result_match.group(1).strip()
+
+        # Parse the result according to the chosen type
+        try:
+            parsed_result = self._convert_to_type(content, chosen_type)
+            return parsed_result, chosen_type
+        except (json.JSONDecodeError, ValueError, TypeError) as e:
+            raise AugmentParseError(
+                f"Could not parse result as {chosen_type.__name__}: {e}"
+            )
+
+    @contextmanager
+    def session(
+        self, session_id: Optional[str] = None
+    ) -> Generator["Agent", None, None]:
+        """
+        Create a session context for maintaining conversation continuity.
+
+        By default, each run() call creates a fresh session. Use this context
+        manager to maintain conversation continuity across multiple calls.
+
+        Usage:
+            agent = Agent()
+
+            # Without session - each call is independent
+            agent.run("Create a function")
+            agent.run("Test it")  # ❌ Won't remember the function
+
+            # With session - calls remember each other
+            with agent.session() as session:
+                session.run("Create a function called add_numbers")
+                session.run("Now test that function")  # ✅ Remembers add_numbers!
+
+        Args:
+            session_id: Optional session ID (currently ignored, auto-generated)
+
+        Yields:
+            Agent: This agent instance
+        """
+        # Mark that we're in a session
+        self._in_session = True
+
+        try:
+            yield self
+        finally:
+            # End the session
+            self._in_session = False
+
+    def _get_type_description(self, return_type: Type) -> str:
+        """Get the JSON structure description for the type."""
+        # Built-in types
+        if return_type in (int, float, str, bool):
+            return "your_result_value"
+        elif return_type in (list, dict):
+            return "your_result_value"
+
+        # Generic types (list[SomeClass], dict[str, SomeClass], etc.)
+        elif hasattr(return_type, "__origin__") or get_origin(return_type) is not None:
+            origin = get_origin(return_type)
+            args = get_args(return_type)
+
+            if origin is list:
+                if args and len(args) == 1:
+                    # list[SomeClass] - create example array with one element
+                    element_type = args[0]
+                    if is_dataclass(element_type):
+                        field_names = [f.name for f in fields(element_type)]
+                        element_example = {
+                            name: f"<{name}_value>" for name in field_names
+                        }
+                        return json.dumps([element_example], indent=2)
+                    elif hasattr(element_type, "__name__") and issubclass(
+                        element_type, Enum
+                    ):
+                        return '["<enum_value>"]'
+                    else:
+                        return '["<element_value>"]'
+                else:
+                    return "your_result_value"  # Plain list
+            elif origin is dict:
+                return "your_result_value"  # For now, treat as plain dict
+            else:
+                return "your_result_value"  # Other generic types
+
+        # Dataclass
+        elif is_dataclass(return_type):
+            field_names = [f.name for f in fields(return_type)]
+            example = {name: f"<{name}_value>" for name in field_names}
+            return json.dumps(example, indent=2)
+
+        # Enum
+        elif issubclass(return_type, Enum):
+            return "your_enum_value"
+
+        else:
+            raise ValueError(f"Unsupported return type: {return_type}")
+
+    def _get_type_instructions(self, return_type: Type) -> str:
+        """Get specific instructions for the type."""
+        # Built-in types
+        if return_type is int:
+            return "IMPORTANT: Put your result inside <augment-agent-result> tags. Return just the integer number (no quotes)."
+        elif return_type is float:
+            return "IMPORTANT: Put your result inside <augment-agent-result> tags. Return just the decimal number (no quotes)."
+        elif return_type is str:
+            return "IMPORTANT: Put your result inside <augment-agent-result> tags. Return the string value in double quotes."
+        elif return_type is bool:
+            return "IMPORTANT: Put your result inside <augment-agent-result> tags. Return either true or false (no quotes)."
+        elif return_type is list:
+            return "IMPORTANT: Put your result inside <augment-agent-result> tags. Return a JSON array: [item1, item2, item3]"
+        elif return_type is dict:
+            return 'IMPORTANT: Put your result inside <augment-agent-result> tags. Return a JSON object: {"key1": "value1", "key2": "value2"}'
+
+        # Generic types (list[SomeClass], dict[str, SomeClass], etc.)
+        elif hasattr(return_type, "__origin__") or get_origin(return_type) is not None:
+            origin = get_origin(return_type)
+            args = get_args(return_type)
+
+            if origin is list:
+                if args and len(args) == 1:
+                    element_type = args[0]
+                    if is_dataclass(element_type):
+                        field_info = []
+                        for field in fields(element_type):
+                            field_info.append(
+                                f"    - {field.name}: {_get_type_name(field.type)}"
+                            )
+
+                        return f"""IMPORTANT: Put your JSON array inside <augment-agent-result> tags.
+
+Return a JSON array of objects, each with these fields:
+{chr(10).join(field_info)}
+
+Example:
+<augment-agent-result>
+[{{"field1": value1, "field2": value2}}, {{"field1": value3, "field2": value4}}]
+</augment-agent-result>"""
+                    elif hasattr(element_type, "__name__") and issubclass(
+                        element_type, Enum
+                    ):
+                        enum_values = [f'"{e.value}"' for e in element_type]
+                        return f"IMPORTANT: Put your result inside <augment-agent-result> tags. Return a JSON array of enum values: {enum_values}"
+                    else:
+                        return f"IMPORTANT: Put your result inside <augment-agent-result> tags. Return a JSON array of {_get_type_name(element_type)} values: [value1, value2, value3]"
+                else:
+                    return "IMPORTANT: Put your result inside <augment-agent-result> tags. Return a JSON array: [item1, item2, item3]"
+            elif origin is dict:
+                return 'IMPORTANT: Put your result inside <augment-agent-result> tags. Return a JSON object: {"key1": "value1", "key2": "value2"}'
+            else:
+                return "IMPORTANT: Put your result inside <augment-agent-result> tags. Return the appropriate JSON structure for your type."
+
+        # Dataclass
+        elif is_dataclass(return_type):
+            field_info = []
+            for field in fields(return_type):
+                field_info.append(f"  - {field.name}: {_get_type_name(field.type)}")
+
+            return f"""IMPORTANT: Put your JSON object inside <augment-agent-result> tags.
+
+Return a JSON object with these exact fields:
+{chr(10).join(field_info)}
+
+Example:
+<augment-agent-result>
+{{"field1": value1, "field2": value2}}
+</augment-agent-result>"""
+
+        # Enum
+        elif issubclass(return_type, Enum):
+            enum_values = [f'"{e.value}"' for e in return_type]
+            return f"IMPORTANT: Put your result inside <augment-agent-result> tags. Return one of these exact values: {', '.join(enum_values)}"
+
+        else:
+            raise ValueError(f"Unsupported return type: {return_type}")
+
+    def _parse_typed_response(self, response: str, return_type: Type[T]) -> T:
+        """Parse the agent's structured response into the desired type."""
+        # Extract message from <augment-agent-message> tags
+        message_match = re.search(
+            r"<augment-agent-message>\s*(.*?)\s*</augment-agent-message>",
+            response,
+            re.DOTALL,
+        )
+        if message_match:
+            self.last_model_answer = message_match.group(1).strip()
+        else:
+            self.last_model_answer = None
+
+        # Extract content from <augment-agent-result> tags
+        result_match = re.search(
+            r"<augment-agent-result>\s*(.*?)\s*</augment-agent-result>",
+            response,
+            re.DOTALL,
+        )
+
+        if not result_match:
+            raise AugmentParseError(
+                "No structured result found. Expected <augment-agent-result> tags in response."
+            )
+
+        content = result_match.group(1).strip()
+
+        try:
+            return self._convert_to_type(content, return_type)
+        except (json.JSONDecodeError, ValueError, TypeError) as e:
+            raise AugmentParseError(
+                f"Could not parse result as {return_type.__name__}: {e}"
+            )
+
+    def _convert_to_type(self, content: str, return_type: Type[T]) -> T:
+        """Convert string content to the specified Python type."""
+
+        # Special case for str - don't JSON parse it, just return as-is
+        if return_type is str:
+            return content  # type: ignore[return-value]
+
+        # Built-in types that need JSON parsing
+        if return_type in (int, float, bool, list, dict):
+            parsed = json.loads(content)
+
+            if not isinstance(parsed, return_type):
+                raise ValueError(
+                    f"Expected {return_type.__name__}, got {type(parsed).__name__}"
+                )
+
+            return parsed
+
+        # Generic types (list[SomeClass], dict[str, SomeClass], etc.)
+        elif hasattr(return_type, "__origin__") or get_origin(return_type) is not None:
+            origin = get_origin(return_type)
+            args = get_args(return_type)
+
+            if origin is list:
+                parsed = json.loads(content)
+                if not isinstance(parsed, list):
+                    raise ValueError(f"Expected list, got {type(parsed).__name__}")
+
+                if args and len(args) == 1:
+                    element_type = args[0]
+                    # Convert each element to the specified type
+                    result = []
+                    for item in parsed:
+                        if is_dataclass(element_type):
+                            if not isinstance(item, dict):
+                                raise ValueError(
+                                    f"Expected dict for dataclass element, got {type(item).__name__}"
+                                )
+                            result.append(element_type(**item))  # type: ignore[misc]
+                        elif hasattr(element_type, "__name__") and issubclass(
+                            element_type, Enum
+                        ):
+                            result.append(element_type(item))
+                        else:
+                            # For basic types, just validate and append
+                            if not isinstance(item, element_type):
+                                raise ValueError(
+                                    f"Expected {element_type.__name__}, got {type(item).__name__}"
+                                )
+                            result.append(item)
+                    return result  # type: ignore[return-value]
+                else:
+                    # Plain list without type parameter
+                    return parsed  # type: ignore[return-value]
+            elif origin is dict:
+                parsed = json.loads(content)
+                if not isinstance(parsed, dict):
+                    raise ValueError(f"Expected dict, got {type(parsed).__name__}")
+                return parsed  # type: ignore[return-value]
+            else:
+                # Other generic types - try basic JSON parsing
+                return json.loads(content)  # type: ignore[return-value,no-any-return]
+
+        # Dataclass
+        elif is_dataclass(return_type):
+            parsed = json.loads(content)
+
+            if not isinstance(parsed, dict):
+                raise ValueError(
+                    f"Expected dict for dataclass, got {type(parsed).__name__}"
+                )
+
+            return return_type(**parsed)
+
+        # Enum
+        elif issubclass(return_type, Enum):
+            # Try JSON parsing first (for quoted values)
+            try:
+                parsed = json.loads(content)
+                return return_type(parsed)  # type: ignore[return-value]
+            except json.JSONDecodeError:
+                # If JSON parsing fails, try direct string value
+                return return_type(content.strip())  # type: ignore[return-value]
+
+        else:
+            raise ValueError(f"Unsupported return type: {return_type}")
+
+    def get_workspace_path(self) -> Path:
+        """
+        Get the workspace path for this agent.
+
+        Returns:
+            Path object representing the workspace root
+        """
+        return self.workspace_path
+
+    @property
+    def session_id(self) -> Optional[str]:
+        """
+        Get the current session ID.
+
+        Returns the session ID only when inside a session context manager.
+        Returns None for standalone run() calls.
+        """
+        if self._in_session and self._acp_client:
+            return self._acp_client.session_id  # type: ignore[attr-defined,no-any-return]
+        return None
+
+    def __repr__(self) -> str:
+        """String representation of the Agent."""
+        if self.session_id:
+            return f"Agent(workspace_path='{self.workspace_path}', session_id='{self.session_id}')"
+        return f"Agent(workspace_path='{self.workspace_path}')"
+
+    def __del__(self) -> None:
+        """Cleanup when agent is destroyed."""
+        # Check if attributes exist (in case __init__ failed)
+        if hasattr(self, "_acp_client") and self._acp_client:
+            try:
+                self._acp_client.stop()
+            except Exception:
+                pass  # Ignore errors during cleanup
+
+    @staticmethod
+    def _validate_workspace_path(workspace_root: Optional[Union[str, Path]]) -> Path:
+        """
+        Validate and resolve workspace path.
+
+        Args:
+            workspace_root: User-provided workspace path or None
+
+        Returns:
+            Resolved Path object
+
+        Raises:
+            AugmentWorkspaceError: If path is invalid
+        """
+        if workspace_root is None:
+            return Path.cwd()
+
+        path = Path(workspace_root).resolve()
+
+        if not path.exists():
+            raise AugmentWorkspaceError(f"Workspace path does not exist: {path}")
+
+        if not path.is_dir():
+            raise AugmentWorkspaceError(f"Workspace path is not a directory: {path}")
+
+        return path
+
+    @staticmethod
+    def get_available_models() -> List[Model]:
+        """
+        Get the list of available AI models for the current account.
+
+        This method calls `auggie model list` to retrieve the available models.
+
+        Returns:
+            List of Model objects containing id, name, and description
+
+        Raises:
+            AugmentNotFoundError: If auggie CLI is not found
+            AugmentCLIError: If the CLI command fails
+
+        Example:
+            >>> models = Agent.get_available_models()
+            >>> for model in models:
+            ...     print(f"{model.name} [{model.id}]")
+            ...     print(f"  {model.description}")
+            Claude Sonnet 4.5 [sonnet4.5]
+              Anthropic Claude Sonnet 4.5, 200k context
+        """
+        # Check if auggie is available
+        try:
+            result = subprocess.run(
+                ["auggie", "model", "list"],
+                capture_output=True,
+                text=True,
+                timeout=30,
+            )
+        except FileNotFoundError:
+            raise AugmentNotFoundError(
+                "auggie CLI not found. Please install auggie and ensure it's in your PATH."
+            )
+        except subprocess.TimeoutExpired:
+            raise AugmentCLIError("Command timed out after 30 seconds", -1, "")
+
+        if result.returncode != 0:
+            raise AugmentCLIError(
+                f"Failed to get model list: {result.stderr}",
+                result.returncode,
+                result.stderr,
+            )
+
+        # Parse the output
+        models = []
+        lines = result.stdout.strip().split("\n")
+
+        i = 0
+        while i < len(lines):
+            line = lines[i].strip()
+
+            # Look for lines that start with " - " (model entries)
+            if line.startswith("- "):
+                # Extract name and id from the format: " - Model Name [model-id]"
+                match = re.match(r"^- (.+?)\s+\[([^\]]+)\]$", line)
+                if match:
+                    name = match.group(1).strip()
+                    model_id = match.group(2).strip()
+
+                    # Next line should be the description (indented)
+                    description = ""
+                    if i + 1 < len(lines):
+                        next_line = lines[i + 1].strip()
+                        if next_line and not next_line.startswith("- "):
+                            description = next_line
+                            i += 1  # Skip the description line
+
+                    models.append(
+                        Model(id=model_id, name=name, description=description)
+                    )
+
+            i += 1
+
+        return models