Skip to content

refactor(tau-coding): reduce coupling in session and TUI hotspots #666

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
darion-yaphet wants to merge 1 commit into
base: main
Choose a base branch
from
Open

refactor(tau-coding): reduce coupling in session and TUI hotspots #666

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
270 changes: 46 additions & 224 deletions src/tau_coding/session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,6 @@
)
from tau_agent.session.entries import SessionEntry
from tau_agent.session.jsonl import entry_to_json_line
from tau_agent.session.tree import SessionTreeError, path_to_entry
from tau_agent.tool_history import ToolHistoryRepair, repair_tool_history
from tau_agent.tools import AgentTool
from tau_agent.types import JSONValue
Expand Down Expand Up @@ -139,6 +138,51 @@
normalize_session_name,
)
from tau_coding.session_stats import SessionStats, calculate_session_stats
from tau_coding.session_terminal import (
TerminalCommandRequest as TerminalCommandRequest,
)
from tau_coding.session_terminal import (
TerminalCommandResult as TerminalCommandResult,
)
from tau_coding.session_terminal import (
parse_terminal_command as parse_terminal_command,
)
from tau_coding.session_terminal import (
terminal_command_context_message,
)
from tau_coding.session_tree import (
SessionTreeBranchResult as SessionTreeBranchResult,
)
from tau_coding.session_tree import (
SessionTreeChoice as SessionTreeChoice,
)
from tau_coding.session_tree import (
detach_missing_parents as _detach_missing_parents,
)
from tau_coding.session_tree import (
is_branchable_tree_entry as _is_branchable_tree_entry,
)
from tau_coding.session_tree import (
is_tool_call_tree_entry as _is_tool_call_tree_entry,
)
from tau_coding.session_tree import (
last_parent_id_from_state as _last_parent_id_from_state,
)
from tau_coding.session_tree import (
latest_leaf_entry as _latest_leaf_entry,
)
from tau_coding.session_tree import (
messages_after_entry_on_active_path as _messages_after_entry_on_active_path,
)
from tau_coding.session_tree import (
ordered_tree_entries as _ordered_tree_entries,
)
from tau_coding.session_tree import (
tree_branch_indents as _tree_branch_indents,
)
from tau_coding.session_tree import (
tree_choice_label as _tree_choice_label,
)
from tau_coding.skills import Skill, expand_skill_command, load_skills_with_diagnostics
from tau_coding.system_prompt import (
BuildSystemPromptOptions,
Expand Down Expand Up @@ -226,43 +270,6 @@ class ModelSelectionResult:
changed: bool


@dataclass(frozen=True, slots=True)
class TerminalCommandResult:
"""Result of an input-bar terminal command."""

command: str
output: str
exit_code: int | None
ok: bool
added_to_context: bool


@dataclass(frozen=True, slots=True)
class SessionTreeChoice:
"""One branchable entry in the active session tree."""

entry_id: str
label: str
active: bool = False
is_tool_call: bool = False


@dataclass(frozen=True, slots=True)
class SessionTreeBranchResult:
"""Result of moving the active session tree leaf."""

message: str
input_prefill: str | None = None


@dataclass(frozen=True, slots=True)
class TerminalCommandRequest:
"""Parsed input-bar terminal command request."""

command: str
add_to_context: bool


@dataclass(frozen=True, slots=True)
class SessionResources:
"""Tau-owned resources loaded around a coding session."""
Expand Down Expand Up @@ -2779,7 +2786,7 @@ async def run_terminal_command(
if add_to_context:
await self._flush_pending_message_writes(context=self._diagnostic_context())
context_message = UserMessage(
content=_terminal_command_context_message(
content=terminal_command_context_message(
normalized_command,
result.text,
)
Expand Down Expand Up @@ -3701,167 +3708,6 @@ def is_retryable_huggingface_route_error(message: AssistantMessage) -> bool:
return False


def _detach_missing_parents(entries: list[SessionEntry]) -> list[SessionEntry]:
"""Return entries with dangling parent pointers detached from external history."""
entry_ids = {entry.id for entry in entries}
return [
entry.model_copy(update={"parent_id": None})
if entry.parent_id is not None and entry.parent_id not in entry_ids
else entry
for entry in entries
]


def _last_parent_id_from_state(state: SessionState) -> str | None:
if state.active_leaf_id is not None:
return state.active_leaf_id
if state.entries:
return state.entries[-1].id
return None


def _latest_leaf_entry(entries: list[SessionEntry]) -> LeafEntry | None:
for entry in reversed(entries):
if isinstance(entry, LeafEntry):
return entry
return None


def _is_branchable_tree_entry(entry: SessionEntry) -> bool:
if entry.type in {"compaction", "branch_summary"}:
return True
if entry.type != "message":
return False
return isinstance(entry.message, UserMessage | AssistantMessage)


def _tree_choice_label(entry: SessionEntry, *, branch_indent: int = 0) -> str:
prefix = " " * branch_indent
return f"{prefix}{_tree_entry_title(entry)}"


def _tree_branch_indents(entries: list[SessionEntry]) -> dict[str, int]:
children_by_parent: dict[str | None, list[str]] = {}
for entry in entries:
if entry.type != "leaf":
children_by_parent.setdefault(entry.parent_id, []).append(entry.id)

sibling_indexes = {
child_id: index
for children in children_by_parent.values()
for index, child_id in enumerate(children)
}
indents: dict[str, int] = {}
for entry in entries:
if entry.type == "leaf":
continue
parent_indent = indents.get(entry.parent_id, 0) if entry.parent_id is not None else 0
sibling_index = sibling_indexes.get(entry.id, 0)
indents[entry.id] = parent_indent + (1 if sibling_index > 0 else 0)
return indents


def _ordered_tree_entries(entries: list[SessionEntry]) -> tuple[SessionEntry, ...]:
children_by_parent: dict[str | None, list[SessionEntry]] = {}
for entry in entries:
if entry.type != "leaf":
children_by_parent.setdefault(entry.parent_id, []).append(entry)

ordered: list[SessionEntry] = []
seen: set[str] = set()
expanded: set[str | None] = set()

def append_descendants(root_parent_id: str | None) -> None:
# Iterative depth-first walk rather than recursion so a long session (a
# deep root-to-leaf entry chain) cannot exceed Python's recursion limit.
# `expanded` also makes a malformed parent cycle terminate instead of
# recursing forever. Emitting a node's direct children before descending,
# and pushing them reversed so the first child is processed next,
# preserves the original traversal order.
stack: list[str | None] = [root_parent_id]
while stack:
parent_id = stack.pop()
if parent_id in expanded:
continue
expanded.add(parent_id)
children = children_by_parent.get(parent_id, [])
for child in children:
if child.id not in seen:
ordered.append(child)
seen.add(child.id)
for child in reversed(children):
stack.append(child.id)

append_descendants(None)
for entry in entries:
if entry.type != "leaf" and entry.id not in seen:
ordered.append(entry)
seen.add(entry.id)
append_descendants(entry.id)
return tuple(ordered)


def _is_tool_call_tree_entry(entry: SessionEntry) -> bool:
return (
entry.type == "message"
and isinstance(entry.message, AssistantMessage)
and bool(entry.message.tool_calls)
)


def _tree_entry_title(entry: SessionEntry) -> str:
match entry.type:
case "message":
message = entry.message
if (
isinstance(message, AssistantMessage)
and message.tool_calls
and not message.text.strip()
):
tool_names = ", ".join(call.name for call in message.tool_calls)
return f"tool call: {tool_names}"
return f"{message.role}: {_message_text_preview(message)}"
case "compaction":
return f"compaction summary: {_short_preview(entry.summary)}"
case "branch_summary":
return f"branch summary: {_short_preview(entry.summary)}"
case _:
return entry.type


def _message_text_preview(message: AgentMessage) -> str:
return _short_preview(message_text(message))


def _short_preview(text: str, *, limit: int = 72) -> str:
normalized = " ".join(text.split())
if len(normalized) <= limit:
return normalized or "(empty)"
return f"{normalized[: limit - 1]}..."


def _messages_after_entry_on_active_path(
entries: list[SessionEntry],
entry_id: str,
active_leaf_id: str | None,
) -> tuple[AgentMessage, ...]:
if active_leaf_id is None:
return ()
try:
active_path = path_to_entry(entries, active_leaf_id)
except SessionTreeError:
return ()
try:
target_index = next(
index for index, entry in enumerate(active_path) if entry.id == entry_id
)
except StopIteration:
return ()
return tuple(
entry.message for entry in active_path[target_index + 1 :] if entry.type == "message"
)


def _storage_path(storage: SessionStorage) -> Path | None:
path = getattr(storage, "path", None)
return path if isinstance(path, Path) else None
Expand Down Expand Up @@ -4278,30 +4124,6 @@ def _fallback_session_name(first_message: str) -> str | None:
return _sanitize_session_name(first_message)


def _terminal_command_context_message(command: str, output: str) -> str:
return (
"Terminal command executed by the user.\n\n"
f"Command:\n```bash\n{command}\n```\n\n"
f"Output:\n```text\n{output}\n```"
)


def parse_terminal_command(text: str) -> TerminalCommandRequest | None:
"""Parse input-bar terminal command syntax."""
stripped = text.strip()
if stripped.startswith("!!"):
command = stripped[2:].strip()
if not command:
return None
return TerminalCommandRequest(command=command, add_to_context=False)
if stripped.startswith("!"):
command = stripped[1:].strip()
if not command:
return None
return TerminalCommandRequest(command=command, add_to_context=True)
return None


def _category_summary(
before: tuple[tuple[object, ...], ...],
after: tuple[tuple[object, ...], ...],
Expand Down
53 changes: 53 additions & 0 deletions src/tau_coding/session_terminal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
"""Input-bar terminal command types and parsing.

This module deliberately owns only the syntax and value objects. Executing a
command and persisting its result remain responsibilities of ``CodingSession``.
"""

from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True, slots=True)
class TerminalCommandResult:
"""Result of an input-bar terminal command."""

command: str
output: str
exit_code: int | None
ok: bool
added_to_context: bool


@dataclass(frozen=True, slots=True)
class TerminalCommandRequest:
"""Parsed input-bar terminal command request."""

command: str
add_to_context: bool


def terminal_command_context_message(command: str, output: str) -> str:
"""Format terminal output for inclusion in agent context."""
return (
"Terminal command executed by the user.\n\n"
f"Command:\n```bash\n{command}\n```\n\n"
f"Output:\n```text\n{output}\n```"
)


def parse_terminal_command(text: str) -> TerminalCommandRequest | None:
"""Parse input-bar terminal command syntax."""
stripped = text.strip()
if stripped.startswith("!!"):
command = stripped[2:].strip()
if not command:
return None
return TerminalCommandRequest(command=command, add_to_context=False)
if stripped.startswith("!"):
command = stripped[1:].strip()
if not command:
return None
return TerminalCommandRequest(command=command, add_to_context=True)
return None
Loading