204 lines
7.3 KiB
Python
204 lines
7.3 KiB
Python
"""New conversations from old ones: spawn with a seed, fork a copy, merge back."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any, cast
|
|
|
|
from claude_agent_sdk import fork_session_via_store, project_key_for_directory
|
|
|
|
from beaver_gateway.backends.transcript import strip_tool_entries, window_entries
|
|
from beaver_gateway.conversations.messaging import Messaging
|
|
from beaver_gateway.conversations.seeds import SEEDS
|
|
from beaver_gateway.conversations.texts import SeedContext
|
|
|
|
if TYPE_CHECKING:
|
|
from beaver_gateway.backends.capture import TurnCapture
|
|
from beaver_gateway.conversations.kinds import Kind
|
|
from beaver_gateway.storage.models import Conversation
|
|
|
|
__all__ = ["ForkResult", "Spawning"]
|
|
|
|
_log = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ForkResult:
|
|
conversation: Conversation
|
|
text: str
|
|
capture: TurnCapture
|
|
|
|
|
|
class Spawning(Messaging):
|
|
async def spawn(
|
|
self,
|
|
*,
|
|
kind: Kind,
|
|
agent: str | None = None,
|
|
seed: str = "clean",
|
|
parent: Conversation | None = None,
|
|
text: str | None = None,
|
|
title: str | None = None,
|
|
window: int | None = None,
|
|
origin: str = "api",
|
|
binding: tuple[str, str] | None = None,
|
|
flags: dict[str, Any] | None = None,
|
|
attachments: list[dict[str, Any]] | None = None,
|
|
) -> Conversation:
|
|
"""Create a conversation in a window and queue its seed.
|
|
|
|
``binding`` reuses a window that already exists instead of asking the
|
|
home frontend for one. Without ``text`` the seed waits in ``flags``
|
|
and opens the first turn, so a fresh window costs nothing until
|
|
someone speaks.
|
|
"""
|
|
if seed not in SEEDS:
|
|
msg = f"unknown seed {seed!r}"
|
|
raise ValueError(msg)
|
|
if seed == "brief" and not text:
|
|
msg = "seed=brief needs text"
|
|
raise ValueError(msg)
|
|
if agent is None and kind == "branch" and parent is not None:
|
|
agent = parent.agent_name
|
|
agent = agent or self.default_agent(kind)
|
|
if agent is None:
|
|
msg = f"no default agent for kind {kind!r}; pass `agent`"
|
|
raise ValueError(msg)
|
|
session_id: str | None = None
|
|
if seed == "copy":
|
|
if parent is None or parent.session_id is None:
|
|
msg = "seed=copy needs a parent with a session"
|
|
raise ValueError(msg)
|
|
session_id = await self._copy_session(
|
|
parent, window=window, strip_tools=False
|
|
)
|
|
conv = await self.create(
|
|
kind=kind,
|
|
agent=agent,
|
|
parent=parent,
|
|
title=title,
|
|
origin=origin,
|
|
session_id=session_id,
|
|
flags=flags,
|
|
)
|
|
if binding is not None:
|
|
await self.bind(conv, frontend=binding[0], external_id=binding[1])
|
|
else:
|
|
await self.materialize(conv)
|
|
ctx = SeedContext(
|
|
kind=kind, seed=seed, agent=agent, parent=parent, text=text, title=title
|
|
)
|
|
if text is None:
|
|
return await self.set_flags(conv, {"seed": seed, "seed_window": window})
|
|
await self._queue.push(
|
|
conversation_id=cast("int", conv.id),
|
|
priority="user",
|
|
origin=f"seed:{seed}" if seed == "brief" else origin,
|
|
text=await self.seed_text(ctx, window=window),
|
|
attachments=attachments,
|
|
)
|
|
self._ensure_worker(cast("int", conv.id))
|
|
return conv
|
|
|
|
async def fork(
|
|
self,
|
|
conv: Conversation,
|
|
prompt: str,
|
|
*,
|
|
window: int | None = None,
|
|
strip_tools: bool = False,
|
|
title: str | None = None,
|
|
agent: str | None = None,
|
|
) -> ForkResult:
|
|
"""Copy the history into a one-off session, run ``prompt`` on it, close it."""
|
|
agent = agent or conv.agent_name
|
|
session_id = await self._copy_session(
|
|
conv, window=window, strip_tools=strip_tools, agent=agent
|
|
)
|
|
child = await self.create(
|
|
kind="fork",
|
|
agent=agent,
|
|
parent=conv,
|
|
title=title or f"fork: {conv.title or conv.external_id}",
|
|
origin="system",
|
|
session_id=session_id,
|
|
)
|
|
try:
|
|
text, capture = await self.run_text_turn(
|
|
child, prompt, origin="fork", tools=False
|
|
)
|
|
finally:
|
|
await self._backend(agent).close(child.external_id)
|
|
child = await self.set_status(child, "closed")
|
|
return ForkResult(conversation=child, text=text, capture=capture)
|
|
|
|
async def merge(self, conv: Conversation) -> ForkResult:
|
|
if conv.parent_id is None:
|
|
msg = "merge needs a parent conversation"
|
|
raise ValueError(msg)
|
|
parent = await self.get_row(conv.parent_id)
|
|
if parent is None:
|
|
msg = "parent conversation vanished"
|
|
raise LookupError(msg)
|
|
result = await self.fork(
|
|
conv,
|
|
self._texts.merge_prompt,
|
|
title=f"merge: {conv.title or conv.external_id}",
|
|
)
|
|
if result.text.strip():
|
|
await self.inject(parent, result.text, urgency="normal", origin="merge")
|
|
await self.set_status(conv, "merged")
|
|
await self.mark_closed(conv)
|
|
self._bus.publish(
|
|
"conversation.merged",
|
|
conversation_id=conv.external_id,
|
|
parent=parent.external_id,
|
|
fork=result.conversation.external_id,
|
|
)
|
|
return result
|
|
|
|
async def _copy_session(
|
|
self,
|
|
conv: Conversation,
|
|
*,
|
|
window: int | None,
|
|
strip_tools: bool,
|
|
agent: str | None = None,
|
|
) -> str:
|
|
if conv.session_id is None:
|
|
msg = f"conversation {conv.external_id} has no session to copy"
|
|
raise ValueError(msg)
|
|
live = self._pool.get(conv.external_id)
|
|
if live is not None and live.dirty:
|
|
msg = f"conversation {conv.external_id} has a mirror gap; not forking"
|
|
raise RuntimeError(msg)
|
|
source = self._claude_agent(conv.agent_name)
|
|
target = self._claude_agent(agent) if agent else source
|
|
forked = await fork_session_via_store(
|
|
self._store, conv.session_id, directory=str(source.cwd)
|
|
)
|
|
source_key = {
|
|
"project_key": project_key_for_directory(str(source.cwd)),
|
|
"session_id": forked.session_id,
|
|
}
|
|
target_key = {
|
|
"project_key": project_key_for_directory(str(target.cwd)),
|
|
"session_id": forked.session_id,
|
|
}
|
|
if window is not None or strip_tools or target_key != source_key:
|
|
entries = await self._store.load(cast("Any", source_key)) or []
|
|
trimmed = window_entries(cast("Any", entries), window=window)
|
|
if strip_tools:
|
|
trimmed = strip_tool_entries(trimmed)
|
|
await self._store.delete(cast("Any", source_key))
|
|
await self._store.append(cast("Any", target_key), cast("Any", trimmed))
|
|
_log.info(
|
|
"forked session %s -> %s (window=%s, strip_tools=%s)",
|
|
conv.session_id,
|
|
forked.session_id,
|
|
window,
|
|
strip_tools,
|
|
)
|
|
return forked.session_id
|