caudate-cli 0.1.0__py3-none-any.whl

core/agentic_loop.py

@@ -0,0 +1,731 @@
+"""Agentic Loop — real-time tool-calling cognitive loop.
+
+The LLM decides which tools to call as it reasons, rather than pre-planning
+a DAG. This is the Claude SDK style of agent loop.
+
+Flow:
+  1. Append user message to history
+  2. Call LLM with tools + history
+  3. If response contains tool_calls: execute each, append tool_results, loop
+  4. If response contains text only: return final text
+
+Memory systems (episodic/semantic) are consulted once at the start of a turn
+and injected as context into the system prompt.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, AsyncIterator, Callable
+
+from core.citations import CitationBlock, Document, parse_citations, render_documents_block
+from core.compaction import ContextCompactor
+from core.files import KIND_IMAGE, FileStore
+from core.hooks import HookEvent, HookManager
+from core.permissions import PermissionManager
+from core.schemas import (
+    Episode,
+    StreamEvent,
+    ThinkingBlock,
+    ToolResultBlock,
+    ToolResultStatus,
+    ToolUseBlock,
+)
+from core.thinking import ThinkingManager
+from execution.executor import Executor
+from llm.provider import LLMProvider
+
+logger = logging.getLogger(__name__)
+
+MAX_ITERATIONS = 25
+MAX_CONSECUTIVE_ERRORS = 3
+
+DEFAULT_SYSTEM_PROMPT = """You are Cognos — a local-first cognitive assistant.
+
+You have access to tools for reading/writing files, running shell commands,
+executing Python, searching the web, and thinking. Use them proactively to
+accomplish the user's goal. When you have enough information, respond with a
+final answer as plain text (no tool call).
+
+Be direct and concise. Do not narrate what you are about to do — just do it.
+Prefer the most specific tool for each subtask (e.g. Read over Bash cat).
+"""
+
+
+class AgenticLoop:
+    """Real-time tool-calling agent loop."""
+
+    def __init__(
+        self,
+        llm: LLMProvider,
+        executor: Executor,
+        episodic=None,
+        semantic=None,
+        working=None,
+        system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+        max_iterations: int = MAX_ITERATIONS,
+        thinking: bool = False,
+        on_thinking=None,
+        permissions: PermissionManager | None = None,
+        hooks: HookManager | None = None,
+        compactor: ContextCompactor | None = None,
+        session_id: str | None = None,
+        system_prompt_provider: "Callable[[str], str] | None" = None,
+        thinking_instruction_wrapper: "Callable[[str], str] | None" = None,
+        file_store: FileStore | None = None,
+    ):
+        self.llm = llm
+        self.executor = executor
+        self.episodic = episodic
+        self.semantic = semantic
+        self.working = working
+        self.system_prompt = system_prompt
+        self.max_iterations = max_iterations
+        self.thinking_manager = ThinkingManager(enabled=thinking)
+        self.on_thinking = on_thinking  # optional callback(list[ThinkingBlock])
+        self.permissions = permissions
+        self.hooks = hooks or HookManager()
+        self.compactor = compactor
+        self.session_id = session_id
+        self.system_prompt_provider = system_prompt_provider
+        self.thinking_instruction_wrapper = thinking_instruction_wrapper
+        self.file_store = file_store
+        self.documents: list[Document] = []
+        self.last_citations: list[CitationBlock] = []
+        self.messages: list[dict[str, Any]] = []
+        self._episodes: list[Episode] = []
+        self._thinking_log: list[ThinkingBlock] = []
+        # Caudate observer — set by CognosAgent if the NN subsystem is
+        # available. Stays None for code paths that build AgenticLoop
+        # directly (subagents, tests).
+        self.caudate = None
+
+    @property
+    def thinking(self) -> bool:
+        return self.thinking_manager.enabled
+
+    @thinking.setter
+    def thinking(self, value: bool) -> None:
+        self.thinking_manager.enabled = value
+
+    # ------------------------------------------------------------------
+    # Public entry points
+    # ------------------------------------------------------------------
+
+    async def run(
+        self,
+        user_message: str,
+        attachments: list[str] | None = None,
+        documents: list[Document] | None = None,
+    ) -> str:
+        """Run the loop to completion and return the final text response.
+
+        `attachments` is a list of file_ids previously registered with
+        the FileStore; they are materialized into the user message.
+
+        `documents` is a list of Document objects the model may cite from;
+        they are injected into the system prompt and their citation markers
+        are parsed out of the response (see self.last_citations).
+        """
+        self.documents = documents or []
+        self.last_citations = []
+        await self.hooks.emit(HookEvent.USER_PROMPT_SUBMIT, {"message": user_message})
+        self._start_turn(user_message, attachments=attachments or [])
+        consecutive_errors = 0
+        # Notify Caudate *first* so her top-K is populated before we
+        # build the tool list (ADR-0007 constrained routing reads
+        # _last_prediction).
+        self._notify_caudate_turn_start(user_message)
+        tools = self._caudate_constrained_tools(self.executor.tool_definitions())
+
+        for iteration in range(self.max_iterations):
+            await self._maybe_compact()
+            self._refresh_system_prompt()
+            try:
+                response = await self.llm.chat(self.messages, tools=tools)
+                consecutive_errors = 0
+            except Exception as e:
+                consecutive_errors += 1
+                logger.error(
+                    f"LLM error at iteration {iteration} "
+                    f"(consecutive={consecutive_errors}): {e}"
+                )
+                if consecutive_errors >= MAX_CONSECUTIVE_ERRORS:
+                    msg = f"(giving up after {consecutive_errors} consecutive LLM errors: {e})"
+                    await self.hooks.emit(HookEvent.STOP, {"text": msg, "error": str(e)})
+                    return msg
+                continue
+
+            # Extract thinking blocks from text
+            content, _ = self._extract_thinking(response.content)
+
+            # Append assistant message (use the stripped content)
+            self._append_assistant_message(content, response.tool_calls)
+
+            if not response.tool_calls:
+                if self.documents:
+                    content, self.last_citations = parse_citations(
+                        content, self.documents,
+                    )
+                await self.hooks.emit(HookEvent.STOP, {"text": content})
+                self._notify_caudate_turn_end()
+                return content
+
+            # Execute tool calls
+            for call in response.tool_calls:
+                result_block = await self._execute_call(call)
+                self._append_tool_result(result_block)
+                self._notify_caudate_tool_use(call.name, was_error=result_block.is_error)
+
+        logger.warning(f"Agentic loop hit max_iterations={self.max_iterations}")
+        self._notify_caudate_turn_end(reward=0.2)  # hit cap → low reward
+        return "(max iterations reached)"
+
+    async def run_streaming(self, user_message: str) -> AsyncIterator[StreamEvent]:
+        """Run the loop streaming tokens/events as they arrive."""
+        self._start_turn(user_message)
+        self._notify_caudate_turn_start(user_message)
+        tools = self._caudate_constrained_tools(self.executor.tool_definitions())
+
+        for iteration in range(self.max_iterations):
+            await self._maybe_compact()
+            self._refresh_system_prompt()
+
+            collected_text = ""
+            collected_calls: list[ToolUseBlock] = []
+            stop_reason: str | None = None
+
+            async for event in self.llm.stream(self.messages, tools=tools):
+                if event.type == "text_delta" and event.delta:
+                    collected_text += event.delta
+                elif event.type == "tool_use_end":
+                    collected_calls.append(ToolUseBlock(
+                        id=event.tool_use_id or "",
+                        name=event.tool_name or "",
+                        input=event.tool_input or {},
+                    ))
+                elif event.type == "message_stop":
+                    stop_reason = event.stop_reason
+                yield event
+
+            # Strip thinking blocks before persisting to history
+            collected_text, _ = self._extract_thinking(collected_text)
+
+            # Append assistant message
+            self._append_assistant_message(collected_text, collected_calls)
+
+            if not collected_calls:
+                self._notify_caudate_turn_end()
+                return
+
+            for call in collected_calls:
+                result_block = await self._execute_call(call)
+                self._append_tool_result(result_block)
+                self._notify_caudate_tool_use(call.name, was_error=result_block.is_error)
+                yield StreamEvent(
+                    type="tool_result",
+                    tool_use_id=call.id,
+                    delta=str(result_block.content)[:500],
+                )
+
+        logger.warning(f"Agentic loop (stream) hit max_iterations={self.max_iterations}")
+
+    def reset(self) -> None:
+        """Clear conversation history (new session)."""
+        self.messages = []
+        self._episodes = []
+
+    async def _maybe_compact(self) -> None:
+        """Run context compaction if enabled and threshold reached."""
+        if self.compactor is None:
+            return
+        if not self.compactor.should_compact(self.messages):
+            return
+        await self.hooks.emit(HookEvent.PRE_COMPACT, {"count": len(self.messages)})
+        self.messages = await self.compactor.compact(self.messages)
+
+    def load_messages(self, messages: list[dict[str, Any]]) -> None:
+        """Restore a conversation (e.g. from a saved session)."""
+        self.messages = list(messages)
+
+    # ------------------------------------------------------------------
+    # Internals
+    # ------------------------------------------------------------------
+
+    def _start_turn(
+        self,
+        user_message: str,
+        attachments: list[str] | None = None,
+    ) -> None:
+        """Ensure system prompt is present and append the user message."""
+        if not self.messages:
+            self.messages.append({
+                "role": "system",
+                "content": self._build_system_prompt(user_message),
+            })
+
+        user_payload = self._build_user_message(user_message, attachments or [])
+        self.messages.append({"role": "user", "content": user_payload})
+
+    def _build_user_message(
+        self,
+        text: str,
+        attachments: list[str],
+    ) -> str | list[dict[str, Any]]:
+        """Render the user turn, materializing attachments.
+
+        Returns a plain string if no attachments (or the store is missing),
+        otherwise a multimodal content-block list with:
+          - extracted text for PDFs / text files (inlined as a text block)
+          - image_url blocks for images
+          - a final text block with the user's actual message
+        """
+        if not attachments or self.file_store is None:
+            return text
+
+        blocks: list[dict[str, Any]] = []
+        for file_id in attachments:
+            rec = self.file_store.get(file_id)
+            if rec is None:
+                blocks.append({
+                    "type": "text",
+                    "text": f"[attachment {file_id} not found]",
+                })
+                continue
+
+            if rec.kind == KIND_IMAGE:
+                content = self.file_store.to_image_content(file_id)
+                if content is not None:
+                    blocks.append({
+                        "type": "text",
+                        "text": f"[image: {rec.filename}]",
+                    })
+                    blocks.append(content)
+                continue
+
+            extracted = self.file_store.extract_text(file_id)
+            if extracted:
+                blocks.append({
+                    "type": "text",
+                    "text": f"--- file: {rec.filename} ({rec.kind}) ---\n{extracted}",
+                })
+            else:
+                blocks.append({
+                    "type": "text",
+                    "text": f"[binary file: {rec.filename} at {rec.path}]",
+                })
+
+        blocks.append({"type": "text", "text": text})
+        return blocks
+
+    def _build_system_prompt(self, user_message: str) -> str:
+        """Compose the full system prompt — base + personality + memory + docs + thinking."""
+        base = self.system_prompt
+        if self.system_prompt_provider is not None:
+            try:
+                base = self.system_prompt_provider(base)
+            except Exception as e:
+                logger.debug(f"system_prompt_provider failed: {e}")
+
+        context = self._memory_context(user_message)
+        if context:
+            base = f"{base}\n\n# Relevant context\n{context}"
+
+        if self.documents:
+            base = f"{base}\n\n{render_documents_block(self.documents)}"
+
+        base = self.thinking_manager.augment_system_prompt(base)
+        if self.thinking_manager.enabled and self.thinking_instruction_wrapper is not None:
+            try:
+                base = self.thinking_instruction_wrapper(base)
+            except Exception as e:
+                logger.debug(f"thinking_instruction_wrapper failed: {e}")
+
+        # Caudate hint injection — only at WHISPER level or above.
+        # Earlier levels stay silent so untrained predictions don't
+        # poison the LLM's prompt.
+        hint = self._caudate_prompt_hint()
+        if hint:
+            base = f"{base}\n\n{hint}"
+        return base
+
+    def _caudate_prompt_hint(self) -> str:
+        """Return a hint block to splice into the system prompt, or empty."""
+        if self.caudate is None or not self.caudate.can_whisper():
+            return ""
+        pred = self.caudate._last_prediction
+        if pred is None:
+            return ""
+        # Don't whisper unless reasonably confident.
+        if pred.tool_confidence < self.caudate.cfg.advisor_min_confidence:
+            return ""
+        level = self.caudate.policy.level.label
+        # Tone scales with trust. Whisper = soft hint. Advisor = stronger.
+        if self.caudate.can_control():
+            preface = (
+                "## Caudate (your trained action-selection net) recommends:"
+            )
+        elif self.caudate.can_advise():
+            preface = (
+                "## Caudate (advisor) suggests, based on prior sessions:"
+            )
+        else:
+            preface = (
+                "## Caudate (whispering — still learning) thinks you might want:"
+            )
+        return (
+            f"{preface}\n"
+            f"  next tool:  {pred.tool}      (confidence {pred.tool_confidence:.2f})\n"
+            f"  routing:    {pred.tier}      (confidence {pred.tier_confidence:.2f})\n"
+            f"  thinking:   {'helpful' if pred.think >= 0.5 else 'not needed'} "
+            f"(p={pred.think:.2f})\n"
+            f"  expected reward: {pred.value:.2f}\n"
+            f"You are free to disagree — Caudate ({level}) has not seen "
+            f"this exact context, only patterns from past turns."
+        )
+
+    def _refresh_system_prompt(self) -> None:
+        """Regenerate the system prompt so dynamic state (mood) stays current."""
+        if self.system_prompt_provider is None:
+            return
+        if not self.messages or self.messages[0].get("role") != "system":
+            return
+        last_user = next(
+            (m.get("content", "") for m in reversed(self.messages) if m.get("role") == "user"),
+            "",
+        )
+        self.messages[0] = {
+            "role": "system",
+            "content": self._build_system_prompt(last_user),
+        }
+
+    def _extract_thinking(self, text: str) -> tuple[str, list[ThinkingBlock]]:
+        """Strip <thinking>…</thinking> blocks out of text.
+
+        Logs them to self._thinking_log and fires the on_thinking callback.
+        Returns (stripped_text, blocks).
+        """
+        if not self.thinking_manager.enabled or not text:
+            return text, []
+        blocks, stripped = self.thinking_manager.parse(text)
+        if blocks:
+            self._thinking_log.extend(blocks)
+            if self.on_thinking:
+                try:
+                    self.on_thinking(blocks)
+                except Exception as e:
+                    logger.debug(f"on_thinking callback failed: {e}")
+        return stripped, blocks
+
+    @property
+    def thinking_log(self) -> list[ThinkingBlock]:
+        return list(self._thinking_log)
+
+    def _memory_context(self, query: str) -> str:
+        """Pull relevant episodes/knowledge into the system prompt."""
+        parts: list[str] = []
+        if self.episodic:
+            try:
+                episodes = self.episodic.recall(query)
+                if episodes:
+                    parts.append("Past episodes:")
+                    for ep in episodes[:3]:
+                        outcome = (
+                            ep.result.status.value
+                            if ep.result else "unknown"
+                        )
+                        parts.append(f"  - {ep.action} -> {outcome}")
+            except Exception as e:
+                logger.debug(f"Episodic recall failed: {e}")
+
+        if self.semantic:
+            try:
+                knowledge = self.semantic.recall(query)
+                if knowledge:
+                    parts.append("Known facts:")
+                    for k in knowledge[:3]:
+                        parts.append(f"  - {k.get('content', '')}")
+            except Exception as e:
+                logger.debug(f"Semantic recall failed: {e}")
+        return "\n".join(parts)
+
+    def _append_assistant_message(
+        self,
+        content: str,
+        tool_calls: list[ToolUseBlock],
+    ) -> None:
+        """Append an assistant message (text + optional tool_calls) to history."""
+        msg: dict[str, Any] = {"role": "assistant", "content": content or ""}
+        if tool_calls:
+            msg["tool_calls"] = [
+                {
+                    "id": tc.id,
+                    "type": "function",
+                    "function": {
+                        "name": tc.name,
+                        "arguments": _json_dump(tc.input),
+                    },
+                }
+                for tc in tool_calls
+            ]
+        self.messages.append(msg)
+
+    def _append_tool_result(self, block: ToolResultBlock) -> None:
+        """Append a tool_result message in LiteLLM/OpenAI format."""
+        self.messages.append({
+            "role": "tool",
+            "tool_call_id": block.tool_use_id,
+            "content": str(block.content),
+        })
+
+    async def _execute_call(self, call: ToolUseBlock) -> ToolResultBlock:
+        """Execute one tool call and return a ToolResultBlock."""
+        logger.info(f"Tool call: {call.name}({call.input})")
+
+        # PRE_TOOL_USE hook — can block or mutate args
+        pre_response = await self.hooks.emit(
+            HookEvent.PRE_TOOL_USE,
+            {"tool": call.name, "tool_name": call.name, "args": call.input},
+        )
+        if pre_response.get("block"):
+            reason = pre_response.get("reason", "blocked by hook")
+            return ToolResultBlock(
+                tool_use_id=call.id,
+                content=f"Blocked by pre_tool_use hook: {reason}",
+                is_error=True,
+            )
+        if "modified_args" in pre_response:
+            call.input = pre_response["modified_args"]
+
+        if self.permissions is not None:
+            allowed, reason = await self.permissions.check(call.name, call.input)
+            if not allowed:
+                return ToolResultBlock(
+                    tool_use_id=call.id,
+                    content=f"Permission denied ({reason})",
+                    is_error=True,
+                )
+
+        result = await self.executor.execute_tool(call.name, call.input)
+
+        hook_event = (
+            HookEvent.POST_TOOL_USE
+            if result.status == ToolResultStatus.SUCCESS
+            else HookEvent.POST_TOOL_USE_FAILURE
+        )
+        await self.hooks.emit(hook_event, {
+            "tool": call.name,
+            "tool_name": call.name,
+            "args": call.input,
+            "status": result.status.value,
+            "output": str(result.output)[:500] if result.output else None,
+            "error": result.error,
+        })
+
+        is_error = result.status != ToolResultStatus.SUCCESS
+        content = result.error if is_error else result.output
+
+        if self.working:
+            try:
+                self.working.store(f"tool_{call.name}_last", content)
+            except Exception:
+                pass
+
+        return ToolResultBlock(
+            tool_use_id=call.id,
+            content=content,
+            is_error=is_error,
+        )
+
+    # ------------------------------------------------------------------
+    # Caudate observer hooks — silent no-ops if caudate is None.
+    # The observer is the bridge between this loop and the NN subsystem;
+    # see nn/observer.py.
+    # ------------------------------------------------------------------
+
+    def _notify_caudate_turn_start(self, user_message: str) -> None:
+        if self.caudate is None:
+            return
+        try:
+            recent = []
+            for m in self.messages[-self.caudate.cfg.msg_window:]:
+                role = m.get("role", "?")
+                c = m.get("content")
+                if isinstance(c, list):
+                    c = " ".join(b.get("text", "") for b in c if isinstance(b, dict))
+                if c:
+                    recent.append(f"{role}: {str(c)[:400]}")
+
+            # Vision channel: scan recent messages for [file:<id>] markers
+            # and resolve any that are images via the FileStore.
+            image_paths = self._collect_recent_image_paths()
+
+            # Hand the executor's current tool registry to Caudate so
+            # the contrastive tool head has real candidates to score.
+            # Without this the head only sees the synthetic <no_tool>
+            # slot and the trained discrimination weights stay dormant.
+            available_tools = self._caudate_available_tools()
+
+            pred = self.caudate.on_turn_start(
+                recent_messages=recent,
+                image_paths=image_paths,
+                available_tools=available_tools,
+            )
+            # CONTROLLER level: Caudate gates thinking. If she predicts
+            # thinking won't help with high confidence, save the tokens.
+            if pred is not None and self.caudate.can_control():
+                if pred.think < 0.3 and self.thinking:
+                    self.thinking = False
+                    logger.info("Caudate disabled thinking for this turn (controller)")
+                elif pred.think > 0.7 and not self.thinking:
+                    self.thinking = True
+                    logger.info("Caudate enabled thinking for this turn (controller)")
+        except Exception as e:
+            logger.debug(f"caudate.on_turn_start failed: {e}")
+
+    def _collect_recent_image_paths(self) -> list[str]:
+        """Pull image-kind file paths off recent messages.
+
+        Two sources:
+          1. Multimodal user content blocks of type 'image_url' — these
+             carry data: URLs we can't pass to CLIP, so we skip them.
+          2. `[file:<uuid>]` markers in either user or assistant text.
+             We resolve each id via the FileStore and keep only ones
+             whose kind is 'image'.
+        """
+        if self.file_store is None:
+            return []
+        import re as _re
+        paths: list[str] = []
+        seen: set[str] = set()
+        marker_re = _re.compile(r"\[file:([a-f0-9-]+)(?:\s+[^\]]+)?\]")
+        for m in self.messages[-self.caudate.cfg.msg_window:] if self.caudate else []:
+            content = m.get("content")
+            if isinstance(content, list):
+                content = " ".join(
+                    b.get("text", "") for b in content if isinstance(b, dict)
+                )
+            if not isinstance(content, str):
+                continue
+            for fid in marker_re.findall(content):
+                if fid in seen:
+                    continue
+                seen.add(fid)
+                rec = self.file_store.get(fid)
+                if rec is not None and rec.kind == "image":
+                    paths.append(rec.path)
+        # Cap to image_window so the encoder doesn't pad uselessly.
+        if self.caudate is not None:
+            paths = paths[-self.caudate.cfg.image_window:]
+        return paths
+
+    def _caudate_constrained_tools(self, all_tools: list) -> list:
+        """ADR-0007: prune the LLM's tool list to Caudate's top-K when she's
+        graduated to ADVISOR+ and is confident about a real tool.
+
+        No-ops (returns ``all_tools`` unchanged) when:
+          - Caudate is absent or below ADVISOR
+          - No prediction this turn
+          - Predicted tool is ``<no_tool>`` or ``<unk>`` (soft-suggest only)
+          - Confidence is below ``advisor_min_confidence``
+          - Top-K is empty (e.g. checkpoint pre-dates top-K capture)
+          - After filtering, fewer than ``advisor_top_k_floor`` tools survive
+            (pathological — bail rather than cripple the LLM)
+        """
+        if self.caudate is None or not self.caudate.can_advise():
+            return all_tools
+        pred = self.caudate._last_prediction
+        if pred is None:
+            return all_tools
+        if pred.tool in ("<no_tool>", "<unk>"):
+            return all_tools
+        if pred.tool_confidence < self.caudate.cfg.advisor_min_confidence:
+            return all_tools
+        if not pred.tool_top_k:
+            return all_tools
+
+        cfg = self.caudate.cfg
+        k = max(cfg.advisor_top_k, cfg.advisor_top_k_floor)
+        keep_names: set[str] = {
+            name for name, _ in pred.tool_top_k[:k]
+            if name not in ("<no_tool>", "<unk>")
+        }
+        keep_names.add(pred.tool)  # predicted tool is always visible
+
+        def _tool_name(t: Any) -> str:
+            if isinstance(t, dict):
+                fn = t.get("function") or {}
+                return fn.get("name") or t.get("name") or ""
+            return getattr(t, "name", "") or ""
+
+        filtered = [t for t in all_tools if _tool_name(t) in keep_names]
+        if len(filtered) < cfg.advisor_top_k_floor:
+            logger.debug(
+                "Caudate constrained-routing bail: only %d/%d tools survived",
+                len(filtered), len(all_tools),
+            )
+            return all_tools
+        logger.info(
+            "Caudate pruned tools %d -> %d (pred=%s conf=%.2f)",
+            len(all_tools), len(filtered), pred.tool, pred.tool_confidence,
+        )
+        return filtered
+
+    def _caudate_available_tools(self) -> list:
+        """Build the candidate-tool list Caudate scores against.
+
+        Each entry is a ``nn.format.ToolDef`` carrying the tool's name
+        plus its human-readable description. The description is what the
+        contrastive head's text embedder ingests, so a vague description
+        hurts discrimination. Importing ToolDef lazily so an import
+        failure on the nn stack doesn't break the agent path.
+        """
+        try:
+            from nn.format import ToolDef
+        except Exception:
+            return []
+        out: list = []
+        for tool in self.executor._tools.values():
+            try:
+                out.append(ToolDef(
+                    name=tool.name,
+                    description=tool.description or "",
+                ))
+            except Exception:
+                continue
+        return out
+
+    def _notify_caudate_tool_use(self, tool_name: str, was_error: bool) -> None:
+        if self.caudate is None:
+            return
+        try:
+            self.caudate.on_tool_use(tool_name, thinking_used=self.thinking)
+        except Exception as e:
+            logger.debug(f"caudate.on_tool_use failed: {e}")
+        # Per-tool reward signal: a failure now reduces this turn's
+        # observed reward when the turn ends.
+        if was_error and self.caudate is not None and self.caudate._pending is not None:
+            self.caudate._pending.tier_used = self.caudate._pending.tier_used  # noop
+            # mark a transient failure so on_turn_end can read it
+            setattr(self.caudate._pending, "_had_error", True)
+
+    def _notify_caudate_turn_end(self, reward: float | None = None) -> None:
+        if self.caudate is None:
+            return
+        try:
+            if reward is None and self.caudate._pending is not None:
+                # Heuristic if no explicit reward: 0.7 default,
+                # 0.3 if any tool errored mid-turn.
+                had_err = getattr(self.caudate._pending, "_had_error", False)
+                reward = 0.3 if had_err else 0.7
+            self.caudate.on_turn_end(reward=reward)
+        except Exception as e:
+            logger.debug(f"caudate.on_turn_end failed: {e}")
+
+
+def _json_dump(obj: Any) -> str:
+    import json
+    try:
+        return json.dumps(obj)
+    except (TypeError, ValueError):
+        return json.dumps(str(obj))