tau-coding-agent 0.1.0__py3-none-any.whl

tau/agent/service.py

@@ -0,0 +1,369 @@
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from tau.agent.types import AgentConfig, AgentContext, AgentPhase, PromptOptions, ContextUsage
+from tau.hooks.service import Hooks
+from tau.hooks.engine import MessageEndEvent, MessageRollbackEvent, SavePointEvent, SettledEvent
+from tau.message.types import AgentMessage, AssistantMessage, TerminalExecutionMessage, LLMMessage, UserMessage, TextContent, ToolMessage
+from tau.message.utils import strip_unusable_trailing_assistant
+from tau.tool.types import ToolInvocation, ToolResult
+
+if TYPE_CHECKING:
+    from tau.engine.service import Engine
+    from tau.session.manager import SessionManager
+    from tau.runtime.service import Runtime
+    from tau.session.compaction import CompactionPreparation
+
+
+
+def _to_llm_messages(messages: list[AgentMessage]) -> list[LLMMessage]:
+    """Convert AgentMessages to LLM-compatible messages.
+
+    TerminalExecutionMessage   → UserMessage (Ran `cmd`\n```output```)
+    CompactionSummaryMessage → UserMessage with summary wrapped in XML tags
+    CustomMessage and other non-LLM types → skipped
+    Empty AssistantMessages are visual-only markers (aborts, persisted API/credit
+    errors) and are skipped — an assistant turn with neither content nor tool
+    calls is invalid to send back and triggers provider 400s.
+    """
+    from tau.message.types import CompactionSummaryMessage, ToolCallContent, ThinkingContent
+    result: list[LLMMessage] = []
+    for msg in messages:
+        if isinstance(msg, CompactionSummaryMessage):
+            text = f"<context-summary>\n{msg.summary}\n</context-summary>"
+            result.append(UserMessage.from_text(text))
+        elif isinstance(msg, TerminalExecutionMessage):
+            if not msg.exclude:
+                result.append(msg.to_user_message())
+        elif isinstance(msg, AssistantMessage):
+            has_usable = any(
+                isinstance(c, (TextContent, ToolCallContent, ThinkingContent))
+                for c in msg.contents
+            )
+            if has_usable:
+                result.append(msg)
+        elif isinstance(msg, (UserMessage, ToolMessage)):
+            result.append(msg)
+    return result
+
+
+class Agent:
+    """
+    High-level agent session tying together Engine and SessionManager.
+
+    Call `invoke()` to run a user turn. The session persists each message
+    and tracks token usage.
+    """
+
+    def __init__(
+        self,
+        engine: Engine,
+        session_manager: SessionManager,
+        config: AgentConfig,
+        hooks: Hooks | None = None,
+    ) -> None:
+        self._engine = engine
+        self._session_manager = session_manager
+        self._config = config
+        self._system_prompt: str = config.system_prompt
+        self._context_tokens: int = 0
+        self._context_window: int = config.context_window
+        self._runtime: Runtime | None = None
+        self.hooks = hooks or Hooks()
+
+        self._phase: AgentPhase = AgentPhase.IDLE
+        self._signal: asyncio.Event = asyncio.Event()
+        self._compaction_failures: int = 0
+        self._engine.options.before_tool_call = self._before_tool_call
+        self._engine.options.after_tool_call = self._after_tool_call
+
+    # -------------------------------------------------------------------------
+    # Public interface
+    # -------------------------------------------------------------------------
+
+    @property
+    def cwd(self) -> Path:
+        """Get the current working directory."""
+        return self._config.cwd
+
+    @property
+    def session_manager(self) -> SessionManager:
+        """Get the session manager instance."""
+        return self._session_manager
+
+    def is_idle(self) -> bool:
+        """Check if the agent is idle (not processing)."""
+        return self._engine.is_idle
+
+    def has_pending_messages(self) -> bool:
+        """Check if there are pending messages in the queue."""
+        return self._engine.has_pending_messages()
+
+    def abort(self) -> None:
+        """Request abort of current operation."""
+        self._signal.set()
+
+    def shutdown(self) -> None:
+        """Shutdown the agent."""
+        self._signal.set()
+
+    def update_context_tokens(self) -> None:
+        """Recalculate context token usage."""
+        from tau.session.compaction import estimate_context_tokens
+        session_ctx = self._session_manager.build_session_context()
+        llm_messages = _to_llm_messages(session_ctx.messages)
+        usage = estimate_context_tokens(llm_messages)
+        self._context_tokens = usage.tokens
+
+    def get_context_usage(self) -> ContextUsage | None:
+        """Get current context token usage and limits."""
+        self.update_context_tokens()
+        percent = (self._context_tokens / self._context_window * 100) if self._context_window else None
+        return ContextUsage(
+            tokens=self._context_tokens,
+            context_window=self._context_window,
+            percent=percent,
+        )
+
+    def get_system_prompt(self) -> str:
+        """Get the system prompt for the agent."""
+        return self._system_prompt
+
+    async def wait_for_idle(self) -> None:
+        """Wait until the agent becomes idle."""
+        await self._engine.wait_for_idle()
+
+    async def new_session(self) -> None:
+        """Create a new session."""
+        if self._runtime is not None:
+            await self._runtime.new_session()
+
+    async def fork(self, entry_id: str) -> None:
+        """Fork a session from a specific entry."""
+        if self._runtime is not None:
+            await self._runtime.fork_session(entry_id)
+
+    async def switch_session(self, session_file: Path) -> None:
+        """Switch to a different session."""
+        if self._runtime is not None:
+            await self._runtime.resume_session(session_file)
+
+    # -------------------------------------------------------------------------
+    # Engine-level tool hooks (pass-through)
+    # -------------------------------------------------------------------------
+
+    async def _before_tool_call(
+        self,
+        invocation: ToolInvocation,
+        signal: asyncio.Event | None,
+    ) -> ToolInvocation | None:
+        return invocation
+
+    async def _after_tool_call(
+        self,
+        invocation: ToolInvocation,
+        result: ToolResult,
+        signal: asyncio.Event | None,
+    ) -> ToolResult | None:
+        return result
+
+    # -------------------------------------------------------------------------
+    # Internal helpers
+    # -------------------------------------------------------------------------
+
+    async def _on_message_end(self, event: MessageEndEvent) -> None:
+        """Persist an incoming message to the session and track token usage."""
+        message = event.message
+        if message is None:
+            return
+        match message:
+            case AssistantMessage():
+                total = message.usage.input_tokens + message.usage.output_tokens
+                if total:
+                    self._context_tokens = total
+                self._session_manager.append_message(message)
+            case ToolMessage():
+                self._session_manager.append_message(message)
+            case _:
+                pass
+
+    async def _on_message_rollback(self, event: "MessageRollbackEvent") -> None:
+        """Retract the last ``event.count`` persisted messages from the session.
+
+        Fired when an interrupted tool turn is dropped: the assistant tool-call
+        message and its tool-result message were already written, so remove them
+        to keep the session consistent with what the engine replays.
+        """
+        for _ in range(event.count):
+            if not self._session_manager.remove_last_message():
+                break
+
+    # -------------------------------------------------------------------------
+    # Compaction
+    # -------------------------------------------------------------------------
+
+    async def compact(self, custom_instructions: str | None = None) -> bool:
+        """Manually trigger context compaction. Returns True if compaction ran."""
+        from tau.session.compaction import prepare_compaction
+        from tau.hooks.engine import CompactionEndEvent
+        entries = self._session_manager.get_branch()
+        preparation = prepare_compaction(entries, self._config.compaction)
+        if preparation is None:
+            return False
+        result, from_extension = await self._run_compaction(preparation, entries, manual=True, custom_instructions=custom_instructions)
+        self._session_manager.append_compaction(
+            summary=result.summary,
+            first_kept_entry_id=result.first_kept_entry_id,
+            tokens_before=result.tokens_before,
+        )
+        self._compaction_failures = 0
+        await self.hooks.emit(CompactionEndEvent(
+            manual=True,
+            tokens_before=result.tokens_before,
+            summary_length=len(result.summary),
+            from_extension=from_extension,
+        ))
+        return True
+
+    async def _check_compaction(self) -> None:
+        """Auto-compact if context usage exceeds the threshold. Circuit-breaks after 3 failures."""
+        from tau.session.compaction import (
+            estimate_context_tokens, should_compact,
+            prepare_compaction,
+        )
+        from tau.hooks.engine import CompactionEndEvent
+
+        if self._compaction_failures >= 3:
+            return
+
+        settings = self._config.compaction
+        if not settings.enabled:
+            return
+
+        entries = self._session_manager.get_branch()
+        session_ctx = self._session_manager.build_session_context()
+        llm_messages = _to_llm_messages(session_ctx.messages)
+
+        usage = estimate_context_tokens(llm_messages)
+        if not should_compact(usage.tokens, self._context_window, settings):
+            return
+
+        preparation = prepare_compaction(entries, settings)
+        if preparation is None:
+            return
+
+        try:
+            result, from_extension = await self._run_compaction(preparation, entries, manual=False)
+            self._session_manager.append_compaction(
+                summary=result.summary,
+                first_kept_entry_id=result.first_kept_entry_id,
+                tokens_before=result.tokens_before,
+            )
+            self._compaction_failures = 0
+            await self.hooks.emit(CompactionEndEvent(
+                manual=False,
+                tokens_before=result.tokens_before,
+                summary_length=len(result.summary),
+                from_extension=from_extension,
+            ))
+        except Exception:
+            self._compaction_failures += 1
+
+    async def _run_compaction(self, preparation: "CompactionPreparation", entries: list, manual: bool, custom_instructions: str | None = None) -> tuple:
+        """Emit before_compaction (allowing interception), then run the default algorithm.
+
+        Returns (CompactionResult, from_extension: bool).
+        Extensions may cancel (raises RuntimeError) or supply a custom CompactionResult.
+        Exceptions in before_compaction handlers are swallowed — first non-error result wins,
+        consistent with error-fallthrough behaviour.
+        """
+        from tau.session.compaction import compact as _compact
+        from tau.hooks.types import BeforeCompactionEvent, BeforeCompactionResult, CompactionStartEvent
+
+        before_results = await self.hooks.emit(BeforeCompactionEvent(
+            preparation=preparation,
+            entries=entries,
+            manual=manual,
+        ))
+
+        for res in before_results:
+            if not isinstance(res, BeforeCompactionResult):
+                continue
+            if res.cancel:
+                raise RuntimeError("Compaction cancelled by extension")
+            if res.compaction is not None:
+                return res.compaction, True
+
+        await self.hooks.emit(CompactionStartEvent(manual=manual))
+        result = await _compact(preparation, self._engine.llm, custom_instructions=custom_instructions)  # type: ignore[arg-type]
+        return result, False
+
+    # -------------------------------------------------------------------------
+    # Core turn entry point
+    # -------------------------------------------------------------------------
+
+    async def invoke(self, text: str, options: PromptOptions | None = None) -> None:
+        """Run one user turn."""
+        if self._phase != AgentPhase.IDLE:
+            raise RuntimeError(f"Agent is busy (phase={self._phase!r}). Wait for the current operation to finish.")
+
+        opts = options or PromptOptions()
+
+        session_ctx = self._session_manager.build_session_context()
+        llm_messages = _to_llm_messages(session_ctx.messages)
+        llm_messages = strip_unusable_trailing_assistant(llm_messages, self._session_manager)
+
+        if opts.images:
+            user_message = UserMessage.with_images(text, list(opts.images))
+        elif opts.audio:
+            user_message = UserMessage.with_audio(text, list(opts.audio))
+        elif opts.video:
+            user_message = UserMessage.with_video(text, list(opts.video))
+        else:
+            user_message = UserMessage.from_text(text)
+        self._session_manager.append_message(user_message, meta=opts.meta)
+        llm_messages.append(user_message)
+
+        ctx = AgentContext(
+            system_prompt=self._system_prompt,
+            messages=llm_messages,
+            tools=self._engine.tools,
+        )
+
+        self._signal = asyncio.Event()
+        self._engine.llm.api.options.signal = self._signal
+
+        self._phase = AgentPhase.TURN
+        try:
+            await self._run(ctx)
+        finally:
+            self._phase = AgentPhase.IDLE
+
+        await self.hooks.emit(SavePointEvent())
+
+        await self._check_compaction()
+
+        if not self._engine.has_pending_messages():
+            await self.hooks.emit(SettledEvent())
+
+    async def _run(self, ctx: AgentContext) -> None:
+        unsubscribe = self.hooks.register(
+            'message_end',
+            lambda event: self._on_message_end(event),
+        )
+        unsubscribe_rollback = self.hooks.register(
+            'message_rollback',
+            lambda event: self._on_message_rollback(event),
+        )
+        try:
+            await self._engine.run(ctx, signal=self._signal)
+        finally:
+            unsubscribe()
+            unsubscribe_rollback()
+
+        error = self._engine.state.error_message
+        if error is not None:
+            raise RuntimeError(f"Agent failed: {error}.")

tau/agent/types.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel
+from tau.message.types import LLMMessage
+from tau.session.types import MessageMeta
+from tau.session.compaction import CompactionSettings, DEFAULT_COMPACTION_SETTINGS
+
+if TYPE_CHECKING:
+    from tau.tool.types import Tool
+
+
+class AgentPhase(str, Enum):
+    """Agent execution phase."""
+    IDLE = "idle"
+    TURN = "turn"
+
+
+@dataclass
+class AgentContext:
+    """Snapshot of everything the LLM receives for one turn."""
+    system_prompt: str
+    messages: list[LLMMessage]
+    tools: list[Tool] = field(default_factory=list)
+
+
+class AgentConfig(BaseModel):
+    """Internal runtime config passed to Agent.__init__."""
+    model_config = {'arbitrary_types_allowed': True}
+
+    cwd: Path
+    system_prompt: str = ""
+    model: Any | None = None
+    context_window: int = 200_000
+    compaction: CompactionSettings = DEFAULT_COMPACTION_SETTINGS
+
+
+
+
+class PromptOptions(BaseModel):
+    """Configuration options for prompt submission."""
+    model_config = {'arbitrary_types_allowed': True}
+
+    meta: MessageMeta | None = None
+    images: list[bytes] = []
+    audio: list[bytes] = []
+    video: list[bytes] = []
+
+
+@dataclass
+class ContextUsage:
+    """Token usage and context window statistics."""
+    tokens: int
+    context_window: int
+    percent: float | None = None
+
+

tau/auth/manager.py

@@ -0,0 +1,247 @@
+from __future__ import annotations
+import os
+import json
+from pathlib import Path
+from typing import List
+from tau.inference.provider.registry import ProviderRegistry
+from tau.inference.provider.oauth import OAuthLoginCallbacks
+from tau.settings.paths import get_auth_path
+from tau.auth.types import AuthCredential, AuthStatus, OAuthCredential, APICredential, AuthType, LockResult
+from tau.auth.storage import AuthStorage, FileAuthStorage, InMemoryAuthStorage
+from tau.utils.secrets import resolve_secret
+
+
+def _get_env_api_key(provider: str) -> str | None:
+    """Get API key for a provider from environment variables."""
+    return os.environ.get(f"{provider.upper()}_API_KEY")
+
+
+class AuthManager:
+    """Credential storage with pluggable backends."""
+
+    def __init__(self, registry: ProviderRegistry, storage: AuthStorage):
+        self.registry = registry
+        self.storage = storage
+        self.runtime_overrides: dict[str, str] = {}
+        self._load_error: Exception | None = None
+        self._errors: list[Exception] = []
+        self.data: dict[str, AuthCredential] = self._load()
+
+    @staticmethod
+    def create(registry: ProviderRegistry, auth_path: Path | None = None) -> "AuthManager":
+        """Create AuthManager with file storage."""
+        path = auth_path or get_auth_path()
+        storage = FileAuthStorage(path)
+        return AuthManager(registry, storage)
+
+    @staticmethod
+    def from_storage(registry: ProviderRegistry, storage: AuthStorage) -> "AuthManager":
+        """Create AuthManager with custom storage."""
+        return AuthManager(registry, storage)
+
+    @staticmethod
+    def in_memory(registry: ProviderRegistry, initial: dict = {}) -> "AuthManager":
+        """Create AuthManager with in-memory storage for testing."""
+        storage = InMemoryAuthStorage()
+        storage.with_lock(lambda _: LockResult(result=None, next=json.dumps(initial, indent=2)))
+        return AuthManager.from_storage(registry, storage)
+
+    def _record_error(self, error: Exception) -> None:
+        """Record an error for later retrieval."""
+        self._errors.append(error)
+
+    def _parse_storage_data(self, content: str | None) -> dict[str, AuthCredential]:
+        """Parse credential data from storage JSON."""
+        if not content:
+            return {}
+        raw_data = json.loads(content)
+        data: dict[str, AuthCredential] = {}
+        for k, v in raw_data.items():
+            cred_type = v.get("type")
+            match cred_type:
+                case AuthType.OAuth:
+                    raw_extra = v.get("extra") or {}
+                    extra = {str(ek): str(ev) for ek, ev in raw_extra.items()} if isinstance(raw_extra, dict) else {}
+                    data[k] = OAuthCredential(
+                        access=v.get("access", ""),
+                        refresh=v.get("refresh", ""),
+                        expires=v.get("expires", 0),
+                        extra=extra,
+                    )
+                case AuthType.ApiKey:
+                    data[k] = APICredential(key=v.get("key", ""))
+        return data
+
+    def _load(self) -> dict[str, AuthCredential]:
+        """Load credentials from storage."""
+        try:
+            result = self.storage.with_lock(lambda current: LockResult(result=current))
+            self._load_error = None
+            return self._parse_storage_data(result.result)
+        except Exception as e:
+            self._load_error = e
+            self._record_error(e)
+            return {}
+
+    @staticmethod
+    def _serialize_credential(credential: AuthCredential) -> dict:
+        """Serialize a credential to storable dict format."""
+        if isinstance(credential, OAuthCredential):
+            return {
+                "type": AuthType.OAuth,
+                "access": credential.access,
+                "refresh": credential.refresh,
+                "expires": credential.expires,
+                "extra": dict(credential.extra),
+            }
+        return {"type": AuthType.ApiKey, "key": credential.key}
+
+    def _persist_provider_change(self, provider: str, credential: AuthCredential | None) -> None:
+        """Persist a credential change to storage."""
+        if self._load_error:
+            return
+
+        def update_fn(current: str | None) -> LockResult:
+            """Update storage data with new credential."""
+            current_data = self._parse_storage_data(current)
+            merged = {k: self._serialize_credential(v) for k, v in current_data.items()}
+            if credential:
+                merged[provider] = self._serialize_credential(credential)
+            else:
+                merged.pop(provider, None)
+            return LockResult(result=None, next=json.dumps(merged, indent=2))
+
+        try:
+            self.storage.with_lock(update_fn)
+        except Exception as e:
+            self._record_error(e)
+
+    def reload(self) -> None:
+        """Reload credentials from storage."""
+        self.data = self._load()
+
+    def get(self, provider: str) -> AuthCredential | None:
+        """Return the stored credential for a provider, or None if not found."""
+        return self.data.get(provider)
+
+    def has(self, provider: str) -> bool:
+        """Check if credentials exist for a provider in storage."""
+        return provider in self.data
+
+    def list(self) -> list[str]:
+        """List all providers with stored credentials."""
+        return list(self.data.keys())
+
+    def set(self, provider: str, credential: AuthCredential) -> None:
+        """Store a credential for a provider and persist to storage."""
+        self.data[provider] = credential
+        self._persist_provider_change(provider=provider, credential=credential)
+
+    def remove(self, provider: str) -> None:
+        """Remove the stored credential for a provider and persist to storage."""
+        self.data.pop(provider, None)
+        self._persist_provider_change(provider=provider, credential=None)
+
+    def set_runtime_api_key(self, provider: str, api_key: str) -> None:
+        """Set a runtime API key override (not persisted)."""
+        self.runtime_overrides[provider] = api_key
+
+    def remove_runtime_api_key(self, provider: str) -> None:
+        """Remove a runtime API key override."""
+        self.runtime_overrides.pop(provider, None)
+
+    def get_auth_status(self, provider: str) -> AuthStatus:
+        """Return auth status without exposing credential values."""
+        if self.has(provider):
+            return AuthStatus(configured=True, source="stored")
+        if provider in self.runtime_overrides:
+            return AuthStatus(configured=True, source="runtime", label="--api-key")
+        env_key = f"{provider.upper()}_API_KEY"
+        if os.environ.get(env_key):
+            return AuthStatus(configured=True, source="env", label=env_key)
+        return AuthStatus(configured=False)
+
+    def drain_errors(self) -> List[Exception]:
+        """Return and clear accumulated errors."""
+        drained = list(self._errors)
+        self._errors.clear()
+        return drained
+
+    async def get_api_key(self, provider: str) -> str | None:
+        """Get an API key for a provider, refreshing OAuth tokens if needed."""
+        # 1. Runtime override
+        if provider in self.runtime_overrides:
+            return resolve_secret(self.runtime_overrides[provider])
+
+        credential = self.get(provider)
+
+        match credential:
+            case APICredential():
+                # The stored key may be a literal, "$ENV_VAR", or "!command";
+                # resolved once and cached (see tau.utils.secrets).
+                return resolve_secret(credential.key)
+            case OAuthCredential():
+                oauth_provider = self.registry.text.get_oauth_provider(provider=provider)
+                if not oauth_provider:
+                    return None
+
+                if oauth_provider.is_expired(credential=credential):
+                    refreshed_credential = await self._refresh_oauth_token_with_lock(provider=provider)
+                    if refreshed_credential:
+                        credential = refreshed_credential
+                    else:
+                        return None
+                return oauth_provider.get_api_key(credential=credential)
+
+        # 2. Environment variable fallback
+        return _get_env_api_key(provider)
+
+    async def _refresh_oauth_token_with_lock(self, provider: str) -> OAuthCredential | None:
+        """Refresh an expired OAuth token with file locking to prevent race conditions."""
+        oauth_provider = self.registry.text.get_oauth_provider(provider=provider)
+        if not oauth_provider:
+            return None
+
+        async def refresh_fn(current: str | None) -> LockResult:
+            """Refresh OAuth token in storage."""
+            current_data = self._parse_storage_data(current)
+            credential = current_data.get(provider)
+
+            if not isinstance(credential, OAuthCredential):
+                return LockResult(result=None)
+
+            # Check if another instance already refreshed
+            if not oauth_provider.is_expired(credential=credential):
+                return LockResult(result=credential)
+
+            try:
+                refreshed_credential = await oauth_provider.refresh_token(credential=credential)
+                if credential.extra:
+                    merged_extra = dict(credential.extra)
+                    merged_extra.update(refreshed_credential.extra)
+                    refreshed_credential.extra = merged_extra
+                current_data[provider] = refreshed_credential
+                self.data = current_data
+                serialized = {k: self._serialize_credential(v) for k, v in current_data.items()}
+                return LockResult(result=refreshed_credential, next=json.dumps(serialized, indent=2))
+            except Exception as e:
+                self._record_error(e)
+                return LockResult(result=None)
+
+        result = await self.storage.with_lock_async(refresh_fn)
+        return result.result
+
+    async def login(self, provider: str, callbacks: OAuthLoginCallbacks):
+        """Perform OAuth login flow for a provider and store the resulting credential."""
+        if oauth_provider := self.registry.text.get_oauth_provider(provider):
+            credential = await oauth_provider.login(callbacks=callbacks)
+            self.data[provider] = credential
+            self._persist_provider_change(provider, credential)
+
+    async def logout(self, provider: str):
+        """Perform OAuth logout for a provider and remove the stored credential."""
+        if oauth_provider := self.registry.text.get_oauth_provider(provider):
+            if credential := self.get(provider):
+                if isinstance(credential, OAuthCredential):
+                    await oauth_provider.logout(credential=credential)
+        self.remove(provider)