diff --git a/packages/kagent/src/kagent/loop.py b/packages/kagent/src/kagent/loop.py index 555d71a..ceba792 100644 --- a/packages/kagent/src/kagent/loop.py +++ b/packages/kagent/src/kagent/loop.py @@ -42,6 +42,37 @@ Receives ``(state, assistant_message)``. Return ``False`` to stop.""" +def _record_turn_end( + state: AgentState, + event: TurnEnd, + *, + run_id: str, + turn_index: int, +) -> TokenUsage | None: + """Record a completed turn in the trace and return token usage. + + Delegates trace-recording responsibility to this helper so that the main + loop stays focused on orchestration (SRP / Delegation principles). + """ + state.trace.append( + TraceEntry.assistant( + event.message, + run_id=run_id, + turn_index=turn_index, + usage=event.message.usage, + ) + ) + for tool_msg in event.tool_results: + state.trace.append( + TraceEntry.tool_result( + tool_msg, + run_id=run_id, + turn_index=turn_index, + ) + ) + return event.message.usage + + async def agent_loop( *, llm: ProviderBase, @@ -148,31 +179,14 @@ async def agent_loop( duration_ms=event.llm_duration_ms, ) - # Record in trace - state.trace.append( - TraceEntry.assistant( - event.message, - run_id=run_id, - turn_index=turn_count, - usage=event.message.usage, - ) + # Record turn in trace (delegated to helper) + turn_usage = _record_turn_end( + state, event, run_id=run_id, turn_index=turn_count ) - for tool_msg in event.tool_results: - state.trace.append( - TraceEntry.tool_result( - tool_msg, - run_id=run_id, - turn_index=turn_count, - ) - ) # Accumulate total usage - if event.message.usage: - total_usage = ( - total_usage + event.message.usage - if total_usage - else event.message.usage - ) + if turn_usage is not None: + total_usage = total_usage + turn_usage if total_usage else turn_usage # Turn end hook turn_duration_ms = (time.perf_counter() - turn_t0) * 1000 diff --git a/packages/kcastle/src/kcastle/castle.py b/packages/kcastle/src/kcastle/castle.py index a206ed5..37d50df 100644 --- a/packages/kcastle/src/kcastle/castle.py +++ b/packages/kcastle/src/kcastle/castle.py @@ -149,47 +149,20 @@ def create( if config is None: config = load_config() - config.home.mkdir(parents=True, exist_ok=True) - config.sessions_dir.mkdir(parents=True, exist_ok=True) - config.skills_dir.mkdir(parents=True, exist_ok=True) + cls._ensure_dirs(config) - project_root = find_project_root(Path.cwd()) - project_skills = project_root / ".agent" / "skills" - skill_manager = SkillManager( - user_skills_dir=config.skills_dir, - project_skills_dir=project_skills, - builtin_skills_dir=Path(__file__).resolve().parent / "skills", + skill_manager = cls._build_skill_manager(config) + provider = create_provider(config.active_provider_config()) + skill_tools = create_builtin_tools(workspace=Path.cwd(), skill_manager=skill_manager) + system_prompt = _build_system_prompt( + config, render_compact_skills(skill_manager.all_skills()) ) - skill_manager.discover() - - provider_config = config.active_provider_config() - provider = create_provider(provider_config) - - all_skills = skill_manager.all_skills() - skill_tools = create_builtin_tools( - workspace=Path.cwd(), - skill_manager=skill_manager, + channels = cls._build_channels( + config, + session_id=session_id, + continue_latest=continue_latest, + daemon=daemon, ) - skill_prompts = render_compact_skills(all_skills) - - system_prompt = _build_system_prompt(config, skill_prompts) - - channels: list[Channel] = [] - if config.cli.enabled and not daemon: - channels.append( - CLIChannel( - session_id=session_id, - continue_latest=continue_latest, - ) - ) - if config.telegram.enabled and config.telegram_token and daemon: - bot_username = config.telegram.options.get("bot_username", "") - channels.append( - TelegramChannel( - token=config.telegram_token, - bot_username=str(bot_username), - ) - ) def agent_factory(trace: Trace) -> Agent: return Agent( @@ -204,11 +177,7 @@ def agent_factory(trace: Trace) -> Agent: sessions_dir=config.sessions_dir, agent_factory=agent_factory, ) - - model_manager = ModelManager( - config=config, - session_manager=session_manager, - ) + model_manager = ModelManager(config=config, session_manager=session_manager) return cls( config=config, @@ -220,6 +189,44 @@ def agent_factory(trace: Trace) -> Agent: skill_tools=skill_tools, ) + @staticmethod + def _ensure_dirs(config: CastleConfig) -> None: + """Create all required application directories.""" + config.home.mkdir(parents=True, exist_ok=True) + config.sessions_dir.mkdir(parents=True, exist_ok=True) + config.skills_dir.mkdir(parents=True, exist_ok=True) + + @staticmethod + def _build_skill_manager(config: CastleConfig) -> SkillManager: + """Create and initialise the skill manager with layered discovery.""" + project_root = find_project_root(Path.cwd()) + skill_manager = SkillManager( + user_skills_dir=config.skills_dir, + project_skills_dir=project_root / ".agent" / "skills", + builtin_skills_dir=Path(__file__).resolve().parent / "skills", + ) + skill_manager.discover() + return skill_manager + + @staticmethod + def _build_channels( + config: CastleConfig, + *, + session_id: str | None, + continue_latest: bool, + daemon: bool, + ) -> list[Channel]: + """Create the configured communication channels.""" + channels: list[Channel] = [] + if config.cli.enabled and not daemon: + channels.append(CLIChannel(session_id=session_id, continue_latest=continue_latest)) + if config.telegram.enabled and config.telegram_token and daemon: + bot_username = config.telegram.options.get("bot_username", "") + channels.append( + TelegramChannel(token=config.telegram_token, bot_username=str(bot_username)) + ) + return channels + async def run(self) -> None: """Start all channels and wait until shutdown.""" if not self._channels: diff --git a/packages/kcastle/src/kcastle/session/session.py b/packages/kcastle/src/kcastle/session/session.py index a462170..a8a3699 100644 --- a/packages/kcastle/src/kcastle/session/session.py +++ b/packages/kcastle/src/kcastle/session/session.py @@ -200,15 +200,13 @@ def create( trace_manager = TraceManager(store=trace_store) trace = trace_manager.create(name=name) - agent: Agent = agent_factory(trace) - logger.info("Created session %s", session_id) - return cls( + return cls._assemble( session_dir=session_dir, meta=meta, - agent=agent, trace=trace, trace_manager=trace_manager, + agent_factory=agent_factory, ) @classmethod @@ -229,9 +227,27 @@ def resume( raise ValueError(f"No trace found in session {meta.id}") trace = trace_manager.load(trace_ids[0]) - agent: Agent = agent_factory(trace) - logger.info("Resumed session %s (%d trace entries)", meta.id, len(trace)) + return cls._assemble( + session_dir=session_dir, + meta=meta, + trace=trace, + trace_manager=trace_manager, + agent_factory=agent_factory, + ) + + @classmethod + def _assemble( + cls, + *, + session_dir: Path, + meta: SessionMeta, + trace: Trace, + trace_manager: TraceManager, + agent_factory: AgentFactory, + ) -> Session: + """Assemble a Session from its components (shared by create and resume).""" + agent: Agent = agent_factory(trace) return cls( session_dir=session_dir, meta=meta,