vibecore 0.3.0b1__py3-none-any.whl → 0.3.1__py3-none-any.whl

vibecore/auth/token_manager.py

@@ -0,0 +1,131 @@
+"""Token management for OAuth authentication."""
+
+import asyncio
+import time
+
+import httpx
+
+from vibecore.auth.config import ANTHROPIC_CONFIG
+from vibecore.auth.models import OAuthCredentials
+from vibecore.auth.storage import SecureAuthStorage
+
+
+class TokenRefreshManager:
+    """Manages OAuth token refresh with automatic renewal."""
+
+    def __init__(self, storage: SecureAuthStorage):
+        """
+        Initialize token refresh manager.
+
+        Args:
+            storage: Secure storage for credentials.
+        """
+        self.storage = storage
+        self.refresh_lock = asyncio.Lock()
+        self.refresh_task: asyncio.Task | None = None
+
+    async def get_valid_token(self) -> str:
+        """
+        Get a valid access token, refreshing if necessary.
+
+        Returns:
+            Valid access token.
+
+        Raises:
+            ValueError: If not authenticated or refresh fails.
+        """
+        auth = await self.storage.load("anthropic")
+
+        if not auth:
+            raise ValueError("Not authenticated")
+
+        if auth.type == "api":  # API key auth
+            return auth.key  # type: ignore
+
+        # OAuth auth - check if token needs refresh
+        now = int(time.time() * 1000)
+        buffer_ms = ANTHROPIC_CONFIG.TOKEN_REFRESH_BUFFER_SECONDS * 1000
+        needs_refresh = not auth.access or auth.expires <= now + buffer_ms  # type: ignore
+
+        if not needs_refresh:
+            return auth.access  # type: ignore
+
+        # Refresh token with lock to prevent concurrent refreshes
+        async with self.refresh_lock:
+            # Re-check after acquiring lock
+            auth = await self.storage.load("anthropic")
+            if auth and auth.type == "oauth":
+                now = int(time.time() * 1000)
+                if auth.access and auth.expires > now + buffer_ms:  # type: ignore
+                    return auth.access  # type: ignore
+
+            # Perform refresh
+            if auth and auth.type == "oauth":
+                return await self._refresh_token(auth.refresh)  # type: ignore
+            else:
+                raise ValueError("Cannot refresh non-OAuth credentials")
+
+    async def _refresh_token(self, refresh_token: str) -> str:
+        """
+        Refresh the access token.
+
+        Args:
+            refresh_token: Refresh token.
+
+        Returns:
+            New access token.
+
+        Raises:
+            httpx.HTTPError: If refresh fails after retries.
+        """
+        last_error: Exception | None = None
+
+        # Retry logic
+        for attempt in range(ANTHROPIC_CONFIG.TOKEN_MAX_RETRY_ATTEMPTS):
+            try:
+                # Exponential backoff for retries
+                if attempt > 0:
+                    delay = ANTHROPIC_CONFIG.TOKEN_RETRY_DELAY_MS * (2 ** (attempt - 1)) / 1000
+                    await asyncio.sleep(delay)
+
+                # Make refresh request
+                async with httpx.AsyncClient() as client:
+                    response = await client.post(
+                        ANTHROPIC_CONFIG.TOKEN_EXCHANGE,
+                        headers={
+                            "Content-Type": "application/json",
+                            "Accept": "application/json",
+                        },
+                        json={
+                            "grant_type": "refresh_token",
+                            "refresh_token": refresh_token,
+                            "client_id": ANTHROPIC_CONFIG.OAUTH_CLIENT_ID,
+                        },
+                        timeout=30.0,
+                    )
+
+                    if response.status_code != 200:
+                        error_text = response.text
+                        raise httpx.HTTPError(f"Token refresh failed: {response.status_code} - {error_text}")
+
+                    tokens_data = response.json()
+
+                # Update stored credentials
+                new_credentials = OAuthCredentials(
+                    type="oauth",
+                    refresh=tokens_data.get("refresh_token", refresh_token),
+                    access=tokens_data["access_token"],
+                    expires=int(time.time() * 1000) + tokens_data["expires_in"] * 1000,
+                )
+
+                await self.storage.save("anthropic", new_credentials)
+                return tokens_data["access_token"]
+
+            except Exception as error:
+                last_error = error
+                print(f"Token refresh attempt {attempt + 1} failed: {error}")
+
+        # All retries failed
+        raise ValueError(
+            f"Token refresh failed after {ANTHROPIC_CONFIG.TOKEN_MAX_RETRY_ATTEMPTS} attempts: {last_error}"
+        )

vibecore/cli.py

@@ -1,5 +1,6 @@
 """Vibecore CLI interface using typer."""
 
+import asyncio
 import logging
 from importlib.metadata import version
 from pathlib import Path
@@ -15,6 +16,10 @@ from vibecore.settings import settings
 
 app = typer.Typer()
 
+# Create auth subcommand group
+auth_app = typer.Typer(help="Manage Anthropic authentication")
+app.add_typer(auth_app, name="auth")
+
 
 def version_callback(value: bool):
     """Handle --version flag."""
@@ -57,9 +62,15 @@ def find_latest_session(project_path: Path | None = None, base_dir: Path | None
     return session_files[0].stem
 
 
-@app.command()
-def run(
-    prompt: str | None = typer.Argument(None, help="Prompt text (requires -p flag)"),
+@app.callback(invoke_without_command=True)
+def main(
+    ctx: typer.Context,
+    prompt: str | None = typer.Option(
+        None,
+        "--prompt",
+        "-p",
+        help="Initial prompt to send to the agent (reads from stdin if -p is used without argument)",
+    ),
     continue_session: bool = typer.Option(
         False,
         "--continue",
@@ -75,7 +86,6 @@ def run(
     print_mode: bool = typer.Option(
         False,
         "--print",
-        "-p",
         help="Print response and exit (useful for pipes)",
     ),
     version: bool | None = typer.Option(
@@ -87,6 +97,10 @@ def run(
     ),
 ):
     """Run the Vibecore TUI application."""
+    # If a subcommand was invoked, don't run the main app
+    if ctx.invoked_subcommand is not None:
+        return
+
     # Set up logging
     logging.basicConfig(
         level="WARNING",
@@ -97,14 +111,14 @@ def run(
     logger.addHandler(TextualHandler())
 
     # Create context
-    ctx = VibecoreContext()
+    vibecore_ctx = VibecoreContext()
 
     # Initialize MCP manager if configured
     mcp_servers = []
     if settings.mcp_servers:
         # Create MCP manager
         mcp_manager = MCPManager(settings.mcp_servers)
-        ctx.mcp_manager = mcp_manager
+        vibecore_ctx.mcp_manager = mcp_manager
 
         # Get the MCP servers from the manager
         mcp_servers = mcp_manager.servers
@@ -125,7 +139,7 @@ def run(
         typer.echo(f"Loading session: {session_to_load}")
 
     # Create app
-    app_instance = VibecoreApp(ctx, agent, session_id=session_to_load, print_mode=print_mode)
+    app_instance = VibecoreApp(vibecore_ctx, agent, session_id=session_to_load, print_mode=print_mode)
 
     if print_mode:
         # Run in print mode
@@ -141,10 +155,85 @@ def run(
         app_instance.run()
 
 
-def main():
+@auth_app.command("login")
+def auth_login(
+    provider: str = typer.Argument("anthropic", help="Authentication provider (currently only 'anthropic')"),
+    api_key: str = typer.Option(None, "--api-key", "-k", help="Use API key instead of OAuth"),
+    mode: str = typer.Option(
+        "max", "--mode", "-m", help="OAuth mode: 'max' for claude.ai, 'console' for console.anthropic.com"
+    ),
+):
+    """Authenticate with Anthropic Pro/Max or API key."""
+    if provider.lower() != "anthropic":
+        typer.echo(f"❌ Provider '{provider}' not supported. Currently only 'anthropic' is supported.")
+        raise typer.Exit(1)
+
+    from vibecore.auth.manager import AnthropicAuthManager
+
+    auth_manager = AnthropicAuthManager()
+
+    if api_key:
+        # API key authentication
+        success = asyncio.run(auth_manager.authenticate_with_api_key(api_key))
+        if not success:
+            raise typer.Exit(1)
+    else:
+        # OAuth Pro/Max authentication
+        success = asyncio.run(auth_manager.authenticate_pro_max(mode))
+        if not success:
+            raise typer.Exit(1)
+
+
+@auth_app.command("logout")
+def auth_logout(
+    provider: str = typer.Argument("anthropic", help="Authentication provider"),
+):
+    """Remove stored authentication."""
+    if provider.lower() != "anthropic":
+        typer.echo(f"❌ Provider '{provider}' not supported. Currently only 'anthropic' is supported.")
+        raise typer.Exit(1)
+
+    from vibecore.auth.manager import AnthropicAuthManager
+
+    auth_manager = AnthropicAuthManager()
+    asyncio.run(auth_manager.logout())
+
+
+@auth_app.command("status")
+def auth_status():
+    """Check authentication status."""
+    from vibecore.auth.manager import AnthropicAuthManager
+
+    auth_manager = AnthropicAuthManager()
+
+    if asyncio.run(auth_manager.is_authenticated()):
+        auth_type = asyncio.run(auth_manager.get_auth_type())
+        if auth_type == "oauth":
+            typer.echo("✅ Authenticated with Anthropic Pro/Max (OAuth)")
+        else:
+            typer.echo("✅ Authenticated with Anthropic API key")
+    else:
+        typer.echo("❌ Not authenticated with Anthropic")
+
+
+@auth_app.command("test")
+def auth_test():
+    """Test authentication by making a simple API call."""
+    from vibecore.auth.manager import AnthropicAuthManager
+
+    auth_manager = AnthropicAuthManager()
+
+    typer.echo("🔍 Testing authentication...")
+    success = asyncio.run(auth_manager.test_connection())
+
+    if not success:
+        raise typer.Exit(1)
+
+
+def cli_main():
     """Entry point for the CLI."""
     app()
 
 
 if __name__ == "__main__":
-    main()
+    cli_main()

vibecore/handlers/stream_handler.py

@@ -15,6 +15,7 @@ from agents import (
     ToolCallOutputItem,
 )
 from openai.types.responses import (
+    ResponseCompletedEvent,
     ResponseFunctionToolCall,
     ResponseOutputItemAddedEvent,
     ResponseOutputItemDoneEvent,
@@ -123,6 +124,7 @@ class AgentStreamHandler:
         """
         match event:
             case RawResponsesStreamEvent(data=data):
+                # log(f"RawResponsesStreamEvent data: {data.type}")
                 match data:
                     case ResponseOutputItemAddedEvent(item=ResponseReasoningItem() as item):
                         reasoning_id = item.id
@@ -180,7 +182,15 @@ class AgentStreamHandler:
                         else:
                             await self.handle_tool_call(tool_name, arguments, call_id)
 
+                    case ResponseCompletedEvent():
+                        # When in agent handoff or stop at tool situations, the tools should be in executing status.
+                        # We find all the executing status tool messages and mark them as success.
+                        for tool_message in self.tool_messages.values():
+                            if tool_message.status == MessageStatus.EXECUTING:
+                                tool_message.status = MessageStatus.SUCCESS
+
             case RunItemStreamEvent(item=item):
+                # log(f"RunItemStreamEvent item: {item.type}")
                 match item:
                     case ToolCallItem():
                         pass
@@ -196,6 +206,7 @@ class AgentStreamHandler:
                         await self.handle_message_complete()
 
             case AgentUpdatedStreamEvent(new_agent=new_agent):
+                # log(f"AgentUpdatedStreamEvent new_agent: {new_agent.name}")
                 await self.message_handler.handle_agent_update(new_agent)
 
     async def handle_task_tool_event(self, tool_name: str, tool_call_id: str, event: StreamEvent) -> None:

vibecore/models/anthropic_auth.py

@@ -0,0 +1,226 @@
+"""Anthropic model with Pro/Max authentication support."""
+
+import logging
+from typing import Any, Literal, overload
+
+import litellm
+from agents.agent_output import AgentOutputSchemaBase
+from agents.handoffs import Handoff
+from agents.items import TResponseInputItem
+from agents.model_settings import ModelSettings
+from agents.models.interface import ModelTracing
+from agents.tool import Tool
+from agents.tracing.span_data import GenerationSpanData
+from agents.tracing.spans import Span
+from openai import AsyncStream
+from openai.types.chat import ChatCompletionChunk
+from openai.types.responses import Response
+
+from vibecore.auth.config import ANTHROPIC_CONFIG
+from vibecore.auth.interceptor import AnthropicRequestInterceptor
+from vibecore.auth.manager import AnthropicAuthManager
+from vibecore.auth.storage import SecureAuthStorage
+from vibecore.models.anthropic import AnthropicModel, _transform_messages_for_cache
+
+logger = logging.getLogger(__name__)
+
+
+class AnthropicProMaxModel(AnthropicModel):
+    """Anthropic model with Pro/Max authentication and Claude Code spoofing."""
+
+    def __init__(self, model_name: str, base_url: str | None = None, api_key: str | None = None, use_auth: bool = True):
+        """
+        Initialize AnthropicProMaxModel.
+
+        Args:
+            model_name: Name of the model.
+            base_url: Optional base URL override.
+            api_key: Optional API key (ignored if Pro/Max auth is active).
+            use_auth: Whether to use Pro/Max authentication.
+        """
+        super().__init__(model_name, base_url, api_key)
+        self.use_auth = use_auth
+        self.auth_manager: AnthropicAuthManager | None = None
+        self.interceptor: AnthropicRequestInterceptor | None = None
+
+        if self.use_auth:
+            self._initialize_auth()
+
+    def _initialize_auth(self) -> None:
+        """Initialize authentication components."""
+        storage = SecureAuthStorage()
+        self.auth_manager = AnthropicAuthManager()
+        self.interceptor = AnthropicRequestInterceptor(storage)
+
+        # Check if authenticated (we'll check async later when needed)
+        logger.info("AnthropicProMaxModel initialized with authentication support")
+
+    async def _inject_system_prompt(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        """
+        Inject Claude Code identity into system messages.
+
+        Args:
+            messages: Original messages.
+
+        Returns:
+            Messages with Claude Code identity injected.
+        """
+        if not self.use_auth or not self.interceptor:
+            return messages
+
+        # Check if using Pro/Max auth
+        storage = SecureAuthStorage()
+        auth = await storage.load("anthropic")
+        if not auth or auth.type != "oauth":
+            return messages  # Only inject for Pro/Max users
+
+        # Find or create system message
+        messages_copy = messages.copy()
+        system_index = next((i for i, msg in enumerate(messages_copy) if msg.get("role") == "system"), None)
+
+        if system_index is not None:
+            # Prepend Claude Code identity to existing system message
+            current_content = messages_copy[system_index].get("content", "")
+            messages_copy[system_index]["content"] = f"{ANTHROPIC_CONFIG.CLAUDE_CODE_IDENTITY}\n\n{current_content}"
+        else:
+            # Add new system message at the beginning
+            messages_copy.insert(0, {"role": "system", "content": ANTHROPIC_CONFIG.CLAUDE_CODE_IDENTITY})
+
+        return messages_copy
+
+    async def _apply_auth_headers(self, headers: dict[str, str]) -> dict[str, str]:
+        """
+        Apply authentication and Claude Code headers.
+
+        Args:
+            headers: Original headers.
+
+        Returns:
+            Modified headers with auth and Claude Code spoofing.
+        """
+        if not self.use_auth or not self.interceptor:
+            return headers
+
+        # Use interceptor to apply auth and Claude Code headers
+        modified_request = await self.interceptor.intercept_request(url="https://api.anthropic.com", headers=headers)
+
+        return modified_request.get("headers", headers)
+
+    @overload
+    async def _fetch_response(
+        self,
+        system_instructions: str | None,
+        input: str | list[TResponseInputItem],
+        model_settings: ModelSettings,
+        tools: list[Tool],
+        output_schema: AgentOutputSchemaBase | None,
+        handoffs: list[Handoff],
+        span: Span[GenerationSpanData],
+        tracing: ModelTracing,
+        stream: Literal[True],
+        prompt: Any | None = None,
+    ) -> tuple[Response, AsyncStream[ChatCompletionChunk]]: ...
+
+    @overload
+    async def _fetch_response(
+        self,
+        system_instructions: str | None,
+        input: str | list[TResponseInputItem],
+        model_settings: ModelSettings,
+        tools: list[Tool],
+        output_schema: AgentOutputSchemaBase | None,
+        handoffs: list[Handoff],
+        span: Span[GenerationSpanData],
+        tracing: ModelTracing,
+        stream: Literal[False],
+        prompt: Any | None = None,
+    ) -> Any: ...
+
+    async def _fetch_response(
+        self,
+        system_instructions: str | None,
+        input: str | list[TResponseInputItem],
+        model_settings: ModelSettings,
+        tools: list[Tool],
+        output_schema: AgentOutputSchemaBase | None,
+        handoffs: list[Handoff],
+        span: Span[GenerationSpanData],
+        tracing: ModelTracing,
+        stream: bool = False,
+        prompt: Any | None = None,
+    ) -> Any | tuple[Response, AsyncStream[ChatCompletionChunk]]:
+        """Override _fetch_response to add auth and Claude Code support."""
+        # Store the original litellm.acompletion function
+        original_acompletion = litellm.acompletion
+
+        async def _intercepting_acompletion(*args, **kwargs):
+            """Intercept litellm.acompletion calls to transform messages and headers."""
+            # Only transform for this Anthropic model
+            if kwargs.get("model") == self.model:
+                if "messages" in kwargs:
+                    messages = kwargs["messages"]
+                    logger.debug(f"Intercepting Anthropic API call with {len(messages)} messages")
+
+                    # Add Claude Code identity to system prompt
+                    messages = await self._inject_system_prompt(messages)
+
+                    # Transform messages to add cache_control
+                    messages = _transform_messages_for_cache(messages)
+                    kwargs["messages"] = messages
+
+                # Apply auth headers if available
+                if self.use_auth and self.interceptor:
+                    # Get existing headers or create new dict
+                    headers = kwargs.get("extra_headers", {})
+
+                    # Apply auth and Claude Code headers
+                    headers = await self._apply_auth_headers(headers)
+
+                    # Update kwargs
+                    kwargs["extra_headers"] = headers
+
+                    # For Pro/Max users, prevent API key from being added
+                    storage = SecureAuthStorage()
+                    auth = await storage.load("anthropic")
+                    if auth and auth.type == "oauth":
+                        # CRITICAL: Set api_key to None to prevent litellm from adding x-api-key header
+                        # when using OAuth authentication
+                        kwargs["api_key"] = None
+
+            # Call the original function with transformed kwargs
+            return await original_acompletion(*args, **kwargs)
+
+        try:
+            # Temporarily replace litellm.acompletion
+            litellm.acompletion = _intercepting_acompletion
+
+            # Call the parent's implementation
+            if stream:
+                return await super()._fetch_response(
+                    system_instructions=system_instructions,
+                    input=input,
+                    model_settings=model_settings,
+                    tools=tools,
+                    output_schema=output_schema,
+                    handoffs=handoffs,
+                    span=span,
+                    tracing=tracing,
+                    stream=True,
+                    prompt=prompt,
+                )
+            else:
+                return await super()._fetch_response(
+                    system_instructions=system_instructions,
+                    input=input,
+                    model_settings=model_settings,
+                    tools=tools,
+                    output_schema=output_schema,
+                    handoffs=handoffs,
+                    span=span,
+                    tracing=tracing,
+                    stream=False,
+                    prompt=prompt,
+                )
+        finally:
+            # Always restore the original function
+            litellm.acompletion = original_acompletion

vibecore/settings.py

@@ -13,6 +13,19 @@ from pydantic_settings import BaseSettings, PydanticBaseSettingsSource, Settings
 from vibecore.models import AnthropicModel
 
 
+class AuthSettings(BaseModel):
+    """Configuration for authentication."""
+
+    use_pro_max: bool = Field(
+        default=False,
+        description="Use Anthropic Pro/Max authentication if available",
+    )
+    auto_refresh: bool = Field(
+        default=True,
+        description="Automatically refresh OAuth tokens",
+    )
+
+
 class SessionSettings(BaseModel):
     """Configuration for session storage."""
 
@@ -124,6 +137,12 @@ class Settings(BaseSettings):
             return None
         return v
 
+    # Authentication configuration
+    auth: AuthSettings = Field(
+        default_factory=AuthSettings,
+        description="Authentication configuration",
+    )
+
     # Session configuration
     session: SessionSettings = Field(
         default_factory=SessionSettings,
@@ -136,17 +155,29 @@ class Settings(BaseSettings):
         description="List of MCP servers to connect to",
     )
 
+    rich_tool_names: list[str] = Field(
+        default_factory=list,
+        description="List of tools to render with RichToolMessage (temporary settings)",
+    )
+
     @property
     def model(self) -> str | Model:
         """Get the configured model.
 
-        Returns an AnthropicModel instance if the model name starts with 'anthropic/',
+        Returns an AnthropicProMaxModel instance if auth is enabled and model is Anthropic,
+        returns an AnthropicModel instance if the model name starts with 'anthropic/',
         returns a OpenAIChatCompletionsModel instance if there is a custom base URL set,
         otherwise returns the model name as a plain string (for OpenAI/LiteLLM models).
         """
         custom_base = "OPENAI_BASE_URL" in os.environ
         if self.default_model.startswith("anthropic/"):
-            return AnthropicModel(self.default_model)
+            # Check if Pro/Max auth should be used
+            if self.auth.use_pro_max:
+                from vibecore.models.anthropic_auth import AnthropicProMaxModel
+
+                return AnthropicProMaxModel(self.default_model, use_auth=True)
+            else:
+                return AnthropicModel(self.default_model)
         elif custom_base:
             openai_provider = MultiProvider().openai_provider
             return OpenAIChatCompletionsModel(self.default_model, openai_provider._get_client())

vibecore/widgets/tool_message_factory.py

@@ -9,12 +9,14 @@ import contextlib
 import json
 from typing import Any
 
+from vibecore.settings import settings
 from vibecore.widgets.messages import MessageStatus
 from vibecore.widgets.tool_messages import (
     BaseToolMessage,
     MCPToolMessage,
     PythonToolMessage,
     ReadToolMessage,
+    RichToolMessage,
     TaskToolMessage,
     TodoWriteToolMessage,
     ToolMessage,
@@ -129,6 +131,12 @@ def create_tool_message(
         else:
             return WebFetchToolMessage(url=url, status=status)
 
+    elif tool_name in settings.rich_tool_names:
+        if output is not None:
+            return RichToolMessage(tool_name=tool_name, arguments=arguments, output=output, status=status)
+        else:
+            return RichToolMessage(tool_name=tool_name, arguments=arguments, status=status)
+
     # Default to generic ToolMessage for all other tools
     else:
         if output is not None: