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

tau/session/types.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+import uuid
+from datetime import datetime
+from pathlib import Path
+from enum import Enum
+from dataclasses import dataclass
+from typing import Any, Literal, Annotated, TYPE_CHECKING, List
+
+from pydantic import BaseModel, Field, ConfigDict
+from tau.inference.types import ThinkingLevel
+
+from tau.message.types import AgentMessage, ImageContent, TextContent
+
+if TYPE_CHECKING:
+    pass
+
+
+def generate_timestamp() -> float:
+    return datetime.now().timestamp()
+
+
+def _generate_id() -> str:
+    return str(uuid.uuid4())[:8]
+
+
+SESSION_VERSION = 3
+
+
+class SessionType(str, Enum):
+    SESSION_HEADER = "session"
+    SESSION_MESSAGE = "message"
+    THINKING_LEVEL_CHANGE = "thinking_level_change"
+    MODEL_CHANGE = "model_change"
+    LABEL = "label"
+    CUSTOM_INFO = "custom"
+    SESSION_INFO = "session_info"
+    CUSTOM_MESSAGE = "custom_message"
+    LEAF = "leaf"
+    COMPACTION = "compaction"
+    BRANCH_SUMMARY = "branch_summary"
+
+
+class BaseSessionEntry(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    id: str = Field(default_factory=_generate_id)
+    timestamp: float = Field(default_factory=generate_timestamp)
+    parent_id: str | None = None
+
+
+class SessionHeader(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    type: Literal[SessionType.SESSION_HEADER] = SessionType.SESSION_HEADER
+    version: int = SESSION_VERSION
+    id: str = Field(default_factory=_generate_id)
+    timestamp: float = Field(default_factory=generate_timestamp)
+    cwd: Path
+    parent_session: Path | None = None
+
+
+class SessionInfoEntry(BaseSessionEntry):
+    type: Literal[SessionType.SESSION_INFO] = SessionType.SESSION_INFO
+    name: str | None = None
+
+
+class MessageAttachment(BaseModel):
+    path: str
+    mime_type: str | None = None
+
+
+class MessageMeta(BaseModel):
+    attachments: list[MessageAttachment] | None = None
+
+
+class MessageEntry(BaseSessionEntry):
+    type: Literal[SessionType.SESSION_MESSAGE] = SessionType.SESSION_MESSAGE
+    message: "AgentMessage"
+    meta: MessageMeta | None = None
+
+
+class ThinkingLevelChangeEntry(BaseSessionEntry):
+    type: Literal[SessionType.THINKING_LEVEL_CHANGE] = SessionType.THINKING_LEVEL_CHANGE
+    thinking_level: ThinkingLevel
+
+
+class ModelChangeEntry(BaseSessionEntry):
+    type: Literal[SessionType.MODEL_CHANGE] = SessionType.MODEL_CHANGE
+    model_id: str
+    provider_id: str
+
+
+class LabelEntry(BaseSessionEntry):
+    type: Literal[SessionType.LABEL] = SessionType.LABEL
+    label: str | None = None
+    target_id: str
+
+
+class LeafEntry(BaseSessionEntry):
+    type: Literal[SessionType.LEAF] = SessionType.LEAF
+    target_id: str | None = None
+
+
+class CustomInfoEntry(BaseSessionEntry):
+    type: Literal[SessionType.CUSTOM_INFO] = SessionType.CUSTOM_INFO
+    custom_type: str
+    data: Any | None = None
+
+
+class CustomMessageEntry(BaseSessionEntry):
+    type: Literal[SessionType.CUSTOM_MESSAGE] = SessionType.CUSTOM_MESSAGE
+    custom_type: str
+    content: List["TextContent | ImageContent"]
+    display: bool = True
+    details: Any | None = None
+
+
+class CompactionEntry(BaseSessionEntry):
+    type: Literal[SessionType.COMPACTION] = SessionType.COMPACTION
+    summary: str
+    first_kept_entry_id: str
+    tokens_before: int
+    details: dict[str, Any] | None = None
+
+
+class BranchSummaryEntry(BaseSessionEntry):
+    type: Literal[SessionType.BRANCH_SUMMARY] = SessionType.BRANCH_SUMMARY
+    from_id: str
+    summary: str
+    details: dict[str, Any] | None = None
+    from_hook: bool = False
+    label: str | None = None
+
+
+SessionEntries = (
+    SessionInfoEntry
+    | MessageEntry
+    | ThinkingLevelChangeEntry
+    | ModelChangeEntry
+    | LabelEntry
+    | LeafEntry
+    | CustomInfoEntry
+    | CustomMessageEntry
+    | CompactionEntry
+    | BranchSummaryEntry
+)
+
+SessionEntry = Annotated[
+    SessionEntries,
+    Field(discriminator="type")
+]
+
+SessionFileEntry = Annotated[
+    SessionHeader | SessionEntries,
+    Field(discriminator="type")
+]
+
+
+class SessionTreeNode(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    entry: SessionEntry
+    children: list[SessionTreeNode] = Field(default_factory=list)
+    label: str | None = None
+    label_timestamp: float | None = None
+
+
+class SessionContext(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    messages: list["AgentMessage"]
+    thinking_level: ThinkingLevel
+    model_id: str | None = None
+    provider_id: str | None = None
+
+
+class SessionInfo(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    path: Path
+    id: str
+    cwd: Path
+    name: str | None = None
+    parent_session: Path | None = None
+    created: datetime
+    modified: datetime
+    message_count: int
+
+
+@dataclass
+class SessionOptions:
+    id: str | None = None
+    parent_session: str | None = None

tau/session/utils.py

@@ -0,0 +1,233 @@
+import re
+import uuid
+from uuid_extensions import uuid7str as _uuid7str
+from typing import Any, Callable
+from datetime import datetime
+from pathlib import Path
+from pydantic import TypeAdapter, ValidationError
+
+from tau.session.types import (
+    SessionEntry, SessionHeader, SessionInfo, MessageEntry, SessionFileEntry, SessionType
+)
+from tau.message.types import AgentMessage, LLMMessage, Role, TextContent, ImageContent
+from tau.settings.paths import get_sessions_dir
+
+def create_session_id() -> str:
+    """Create a new session ID using UUIDv7."""
+    return _uuid7str()
+
+def generate_id(by_id: Any) -> str:
+    """
+    Generate a unique short ID (8 hex chars, collision-checked).
+    
+    Args:
+        by_id: A container (like a set or dict) that supports the 'in' operator 
+               to check for existing IDs.
+    """
+    for _ in range(100):
+        new_id = str(uuid.uuid4())[:8]
+        if new_id not in by_id:
+            return new_id
+
+    # Fallback to full UUID if somehow we have collisions
+    return str(uuid.uuid4())
+
+def generate_timestamp() -> float:
+    """Generate a Unix timestamp for the current moment."""
+    now = datetime.now()
+    return now.timestamp()
+
+def get_default_session_dir(cwd: str | Path, sessions_dir: Path | None = None) -> Path:
+    """Return the per-project session directory under ~/.tau/sessions/<encoded-cwd>/."""
+    base = sessions_dir if sessions_dir is not None else get_sessions_dir()
+    resolved = str(Path(cwd).resolve())
+    # Encode the absolute path into a safe directory name: --home-user-project--
+    safe = "--" + re.sub(r"^[/\\]", "", resolved).replace("/", "-").replace("\\", "-").replace(":", "-") + "--"
+    session_dir = base / safe
+    session_dir.mkdir(parents=True, exist_ok=True)
+    return session_dir
+
+def read_session_file(session_file: Path) -> list[SessionFileEntry]:
+    """Load and parse a session file, returning a list of entries."""
+    if not session_file.exists():
+        return []
+
+    adapter = TypeAdapter(SessionFileEntry)
+
+    content = session_file.read_text(encoding="utf-8")
+    entries: list[SessionFileEntry] = []
+
+    for line in content.splitlines():
+        if not line.strip():
+            continue
+        try:
+            entry = adapter.validate_json(line)
+            entries.append(entry)
+        except Exception:
+            continue
+
+    if len(entries) == 0:
+        return []
+
+    header = entries[0]
+
+    if header.type != SessionType.SESSION_HEADER:
+        return []
+
+    return entries
+
+def is_valid_session_file(session_file: Path | str) -> bool:
+    """Check if a file is a valid session file by validating its header."""
+    try:
+        path = Path(session_file)
+        if not path.exists():
+            return False
+
+        with path.open("r", encoding="utf-8") as file:
+            first_line = file.readline().strip()
+
+        if not first_line:
+            return False
+
+        SessionHeader.model_validate_json(first_line)
+        return True
+    except (OSError, ValidationError, ValueError):
+        return False
+
+def find_most_recent_session(session_dir: Path | str) -> Path | None:
+    """Find the most recently modified session file in a directory."""
+    session_dir = Path(session_dir)
+    if not session_dir.is_dir():
+        return None
+
+    candidate_sessions = [p for p in session_dir.glob("*.jsonl") if is_valid_session_file(p)]
+
+    if not candidate_sessions:
+        return None
+
+    most_recent = max(candidate_sessions, key=lambda x: x.stat().st_mtime)
+    return most_recent
+
+def is_message_with_contents(message: AgentMessage) -> bool:
+    """Check if a message is an LLM message with user or assistant role and content."""
+    if not isinstance(message, LLMMessage):
+        return False
+    if message.role not in (Role.USER, Role.ASSISTANT):
+        return False
+    return any(isinstance(c, (TextContent, ImageContent)) for c in message.contents)
+
+def get_last_activity_time(entries: list[SessionEntry]) -> float | None:
+    """Extract the most recent message timestamp from a list of session entries."""
+    last_activity_time = None
+
+    for entry in entries:
+        if not isinstance(entry, MessageEntry):
+            continue
+
+        if not is_message_with_contents(entry.message):
+            continue
+        
+        message_timestamp = getattr(entry.message, "timestamp", None)
+        if message_timestamp is None:
+            timestamp = entry.timestamp
+        elif isinstance(message_timestamp, (int, float)):
+            timestamp = float(message_timestamp)
+        else:
+            timestamp = float(message_timestamp.timestamp())
+
+        last_activity_time = max(last_activity_time or 0.0, timestamp)
+
+    return last_activity_time
+
+def get_session_modified_date(entries: list[SessionEntry], header: SessionHeader | None = None) -> datetime:
+    """Get the modified timestamp of a session, based on last activity or header creation time."""
+    if last_activity_time := get_last_activity_time(entries=entries):
+        return datetime.fromtimestamp(last_activity_time)
+
+    header = header or entries[0]
+    return datetime.fromtimestamp(header.timestamp)
+
+def build_session_info(file: Path) -> SessionInfo | None:
+    """Parse a session file and extract metadata into a SessionInfo object."""
+    content = file.read_text(encoding="utf-8")
+
+    file_entries: list[SessionFileEntry] = []
+    lines = content.strip().splitlines()
+    adapter = TypeAdapter(SessionFileEntry)
+
+    for line in lines:
+        if not line.strip():
+            continue
+        try:
+            file_entries.append(adapter.validate_json(line))
+        except Exception:
+            pass
+
+    if len(file_entries) == 0:
+        return None
+
+    header: SessionHeader | None = None
+    entries: list[SessionEntry] = []
+    message_count = 0
+    for entry in file_entries:
+        if isinstance(entry, SessionHeader):
+            header = entry
+        else:
+            entries.append(entry)
+            if isinstance(entry, MessageEntry):
+                message_count += 1
+
+    if header is None:
+        return None
+    
+    cwd = header.cwd
+    parent_session = header.parent_session
+    created = datetime.fromtimestamp(header.timestamp)
+    modified = get_session_modified_date(entries, header)
+
+    return SessionInfo(
+        path=file,
+        id=header.id,
+        cwd=cwd,
+        parent_session=parent_session,
+        created=created,
+        modified=modified,
+        message_count=message_count
+    )
+
+def list_sessions_from_dir(
+    dir_path: Path | str,
+    on_progress: Callable[[int, int], None] | None = None,
+    progress_offset: int = 0,
+    progress_total: int | None = None
+) -> list[SessionInfo]:
+    """
+    Read all .jsonl session files in a directory and return a list of SessionInfo objects.
+    Optionally reports progress through the on_progress callback.
+    """
+    sessions: list[SessionInfo] = []
+    dir_path = Path(dir_path)
+    
+    if not dir_path.exists() or not dir_path.is_dir():
+        return sessions
+
+    try:
+        files = list(dir_path.glob("*.jsonl"))
+        total = progress_total if progress_total is not None else len(files)
+        loaded = 0
+
+        # We process files sequentially since Python I/O blocking is usually fine here,
+        # but could be updated to use ThreadPoolExecutor if concurrency is strictly needed.
+        for file in files:
+            info = build_session_info(file)
+            loaded += 1
+            if on_progress:
+                on_progress(progress_offset + loaded, total)
+            
+            if info is not None:
+                sessions.append(info)
+                
+    except Exception:
+        pass  # Return what we have on error, or an empty list if early
+
+    return sessions