diff --git a/anton/README.md b/anton/README.md index 776bf6a9..31ac36e3 100644 --- a/anton/README.md +++ b/anton/README.md @@ -305,7 +305,7 @@ When scratchpads are active, relevant lessons are appended to the scratchpad too scratchpad_tool["description"] += f"\n\nLessons from past sessions:\n{wisdom}" ``` -This combines all "when" rules + lessons with `scratchpad-*` topics from both scopes. The content comes from `cortex.get_scratchpad_context()`, which calls `recall_scratchpad_wisdom()` on both hippocampi. +This combines scratchpad-related "when" rules + scratchpad-related lessons from both scopes, ordered by confidence tier (high, then medium/unset, then low) and recency within a tier, trimmed to a token budget (default 2000/scope, ~4 chars/token). The content comes from `cortex.get_scratchpad_context()`, which calls `recall_scratchpad_wisdom()` on both hippocampi. ## The `memorize` Tool @@ -572,7 +572,7 @@ When errored cells exist, they accumulate in a buffer across the turn. At end-of 1. The buffered cells get formatted into a compact post-mortem prompt 2. One LLM call via `LLMClient.generate_object_code` (the cheap coding model) returns a `_DiffPassResult` Pydantic model with extracted lessons 3. Each lesson is wrapped as an `Engram` with `kind="lesson"`, `topic="scratchpad"`, `source="consolidation"`, and routed through `Cortex.encode()` — the same path manual lessons and the consolidator already use -4. Future scratchpad cells see those lessons via the existing `recall_scratchpad_wisdom()` injection into the scratchpad tool description +4. Future scratchpad cells see those lessons via the existing `recall_scratchpad_wisdom()` injection into the scratchpad tool description (confidence tier + recency ordered, budget-limited — not all lessons are guaranteed to survive the cut) The cerebellum is a **producer** only — it generates new lesson entries for the existing storage and retrieval pipeline. There's no parallel storage system, no separate `corrections.md` file. Whatever the consolidator and `/memorize` write to, the cerebellum also writes to. @@ -811,7 +811,7 @@ The Hippocampus handles one scope (global OR project) and is the canonical file- | `recall_rules()` | `rules.md` | Basal Ganglia + OFC | | `recall_lessons(token_budget)` | `lessons.md` (budget-limited, most recent first) | Anterior Temporal Lobe | | `recall_topic(slug)` | `topics/{slug}.md` | Cortical Association Areas | -| `recall_scratchpad_wisdom()` | "when" rules + scratchpad-related lessons + `topics/scratchpad-*.md` | Procedural memory | +| `recall_scratchpad_wisdom(token_budget)` | scratchpad-related "when" rules + scratchpad-related lessons (confidence tier + recency ordered, budget-limited) | Procedural memory | **Encoding methods:** | Method | Writes | Behavior | diff --git a/anton/core/memory/base.py b/anton/core/memory/base.py index 4a5cbf47..ea171c8a 100644 --- a/anton/core/memory/base.py +++ b/anton/core/memory/base.py @@ -53,16 +53,21 @@ def recall_rules(self) -> str: """Return behavioral gates (rules.md equivalent).""" ... - def recall_lessons(self, token_budget: int = 1000) -> str: - """Return semantic facts within the given token budget.""" + def recall_lessons(self, token_budget: int = 1000, exclude_scratchpad: bool = False) -> str: + """Return semantic facts within the given token budget. + + exclude_scratchpad drops entries whose text/topic mentions + "scratchpad" — used by build_memory_context() to avoid double-billing + entries already injected via recall_scratchpad_wisdom(). + """ ... def recall_topic(self, slug: str) -> str: """Return lessons tagged with the given topic slug.""" ... - def recall_scratchpad_wisdom(self) -> str: - """Return procedural knowledge relevant to scratchpad execution.""" + def recall_scratchpad_wisdom(self, token_budget: int = 2000) -> str: + """Return procedural knowledge relevant to scratchpad execution, within budget.""" ... # --- read (Engrams, for inspection / CRUD UI) --- @@ -71,12 +76,22 @@ def get_identities(self) -> list[Engram]: """Return identity entries as Engrams.""" ... - def get_rules(self) -> list[Engram]: - """Return behavioral rules as Engrams.""" + def get_rules(self, exclude_scratchpad_when: bool = False) -> list[Engram]: + """Return behavioral rules as Engrams. + + exclude_scratchpad_when drops "when" rules whose text/topic mentions + "scratchpad" — used by build_memory_context() to avoid double-billing + entries already injected via recall_scratchpad_wisdom(). always/never + rules are never affected. + """ ... - def get_lessons(self, token_budget: int = None) -> list[Engram]: - """Return semantic facts as Engrams, optionally budget-limited.""" + def get_lessons(self, token_budget: int = None, exclude_scratchpad: bool = False) -> list[Engram]: + """Return semantic facts as Engrams, optionally budget-limited. + + exclude_scratchpad drops entries whose text/topic mentions + "scratchpad", applied before the budget cut. + """ ... # --- write --- diff --git a/anton/core/memory/cortex.py b/anton/core/memory/cortex.py index fa8df0ed..4d0b436f 100644 --- a/anton/core/memory/cortex.py +++ b/anton/core/memory/cortex.py @@ -185,8 +185,11 @@ async def build_memory_context(self, user_message: str = "") -> str: if identity: sections.append(f"## Your Memory — Identity\n{identity}") - # 2. Global rules (with smart retrieval) - global_engrams = self.global_hc.get_rules() + # 2. Global rules (with smart retrieval). get_rules(exclude_scratchpad_when=True) + # drops scratchpad-related "when" rules — those are already injected + # into the scratchpad tool description by get_scratchpad_context(), + # and showing them here too would double their token cost. + global_engrams = self.global_hc.get_rules(exclude_scratchpad_when=True) if global_engrams: global_engrams = await self._retrieve_relevant_rules(global_engrams, user_message) if global_engrams: @@ -194,8 +197,8 @@ async def build_memory_context(self, user_message: str = "") -> str: f"## Your Memory — Global Rules\n{self._format_rules_engrams(global_engrams)}" ) - # 3. Project rules (with smart retrieval) - project_engrams = self.project_hc.get_rules() + # 3. Project rules (with smart retrieval) — same scratchpad exclusion. + project_engrams = self.project_hc.get_rules(exclude_scratchpad_when=True) if project_engrams: project_engrams = await self._retrieve_relevant_rules(project_engrams, user_message) if project_engrams: @@ -205,17 +208,18 @@ async def build_memory_context(self, user_message: str = "") -> str: for engram in project_engrams: self._log_read_engram(engram) - # 4. Global lessons - global_lessons = self.global_hc.recall_lessons(token_budget=1000) + # 4. Global lessons. recall_lessons() excludes scratchpad-related + # entries internally — same reasoning as the rules exclusion above. + global_lessons = self.global_hc.recall_lessons(token_budget=1000, exclude_scratchpad=True) if global_lessons: sections.append(f"## Your Memory — Global Lessons\n{global_lessons}") - # 5. Project lessons - project_lessons = self.project_hc.recall_lessons(token_budget=1000) + # 5. Project lessons — same scratchpad exclusion. + project_lessons = self.project_hc.recall_lessons(token_budget=1000, exclude_scratchpad=True) if project_lessons: sections.append(f"## Your Memory — Project Lessons\n{project_lessons}") if self._episodic is not None: - for engram in self.project_hc.get_lessons(token_budget=1000): + for engram in self.project_hc.get_lessons(token_budget=1000, exclude_scratchpad=True): self._log_read_engram(engram) # 6. Minds datasource context (auto-loaded if present) diff --git a/anton/core/memory/hippocampus.py b/anton/core/memory/hippocampus.py index 13790aae..95b4916c 100644 --- a/anton/core/memory/hippocampus.py +++ b/anton/core/memory/hippocampus.py @@ -67,6 +67,39 @@ def _extract_metadata(text: str) -> tuple[str, dict]: return clean_text, meta +def _filter_by_token_budget(engrams: list[Engram], token_budget: int) -> list[Engram]: + """Keep entries, in the given order, until token_budget is spent. + + Callers must pass entries already sorted by priority (most important + first) — this stops at the first entry that would exceed the budget, it + does not repack for a tighter fit. Budget enforced at ~4 chars/token. + """ + char_budget = token_budget * 4 + used = 0 + kept: list[Engram] = [] + for engram in engrams: + length = len(engram.text) + 1 + if used + length > char_budget: + break + kept.append(engram) + used += length + return kept + + +def _is_scratchpad_related(engram: Engram) -> bool: + """Whether an engram's text or topic mentions "scratchpad". + + Used both to select entries INTO recall_scratchpad_wisdom() (the + scratchpad tool description) and to exclude the same entries FROM + get_rules()/get_lessons() when building the system prompt — keeping one + definition avoids the two call sites drifting apart and double-injecting + (or double-excluding) an entry. + """ + if "scratchpad" in engram.text.lower(): + return True + return bool(engram.topic and "scratchpad" in engram.topic.lower()) + + class Hippocampus: """Reads and writes memory traces at a single scope (global OR project). @@ -191,16 +224,16 @@ def clear_identity(self) -> None: # --------- lessons -------------- - def recall_lessons(self, token_budget: int|None = 1000) -> str: + def recall_lessons(self, token_budget: int|None = 1000, exclude_scratchpad: bool = False) -> str: """Load semantic knowledge (lessons.md), most recent first, within budget. Brain analog: Anterior Temporal Lobe — the convergence hub for semantic facts distilled from many episodes. Budget enforced at ~4 chars/token. """ - return self._lessons_to_text(self.get_lessons(token_budget)) + return self._lessons_to_text(self.get_lessons(token_budget, exclude_scratchpad=exclude_scratchpad)) - def get_lessons(self, token_budget: int = None) -> list[Engram]: + def get_lessons(self, token_budget: int = None, exclude_scratchpad: bool = False) -> list[Engram]: """Load semantic knowledge (lessons.md) as Engrams. When token_budget is None (default) returns all entries in file order. @@ -233,24 +266,16 @@ def get_lessons(self, token_budget: int = None) -> list[Engram]: text, meta = _extract_metadata(text) entries.append(Engram(text=text.removeprefix("- "), **meta)) + if exclude_scratchpad: + entries = [e for e in entries if not _is_scratchpad_related(e)] + if token_budget is None: return entries # Reverse entries so most recent are first entries.reverse() - # Budget: ~4 chars per token - char_budget = token_budget * 4 - result_lines = [] - used = sum(len(ln) for ln in result_lines) - - for ln in entries: - if used + len(ln.text) + 1 > char_budget: - break - result_lines.append(ln) - used += len(ln.text) + 1 - - return result_lines + return _filter_by_token_budget(entries, token_budget) def del_lesson(self, id): entries = self.get_lessons() @@ -320,36 +345,46 @@ def recall_topic(self, slug: str) -> str: return self._lessons_to_text(items, header=slug) - def recall_scratchpad_wisdom(self) -> str: + def recall_scratchpad_wisdom(self, token_budget: int = 2000) -> str: """Retrieve procedural knowledge relevant to scratchpad execution. - Returns all "when" rules + lessons whose text contains "scratchpad". - Injected into tool descriptions so the LLM sees them when composing code. + Gathers "when" rules + lessons whose text or topic mentions + "scratchpad", ordered by confidence tier (high, then medium/unset, + then low) and by recency within a tier, then trimmed to token_budget + (~4 chars/token — same convention as get_lessons). Injected into tool + descriptions so the LLM sees them when composing code. """ - parts: list[str] = [] - - # Extract "when" rules - rules = self.recall_rules() - if rules: - in_when = False - for line in rules.splitlines(): - if line.strip().startswith("## When"): - in_when = True - continue - elif line.strip().startswith("## "): - in_when = False - continue - if in_when and line.strip().startswith("- "): - parts.append(line.strip()) + candidates: list[Engram] = [] + + for rule in self.get_rules(): + if rule.kind == "when" and _is_scratchpad_related(rule): + candidates.append(rule) - # Extract scratchpad-related lessons for lesson in self.get_lessons(): - if "scratchpad" in lesson.text.lower() or (lesson.topic and "scratchpad" in lesson.topic.lower()): - entry = f"- {lesson.text}" - if entry not in parts: - parts.append(entry) + if _is_scratchpad_related(lesson): + candidates.append(lesson) + + confidence_rank = {"high": 0, "medium": 1, "low": 2} + + def sort_key(engram: Engram) -> tuple[int, float]: + rank = confidence_rank.get(engram.confidence, 1) + recency = -engram.updated_at.timestamp() if engram.updated_at else float("inf") + return (rank, recency) + + candidates.sort(key=sort_key) - return "\n".join(parts) + # Dedup by text, keeping the higher-priority copy from the sort above + seen_texts: set[str] = set() + deduped: list[Engram] = [] + for engram in candidates: + if engram.text in seen_texts: + continue + seen_texts.add(engram.text) + deduped.append(engram) + + kept = _filter_by_token_budget(deduped, token_budget) + + return "\n".join(f"- {engram.text}" for engram in kept) # ---------- rules ------------------- @@ -367,8 +402,15 @@ def recall_rules(self) -> str: except (OSError, UnicodeDecodeError): return "" - def get_rules(self) -> list[Engram]: - """Load behavioral rules (rules.md) as Engrams, grouped by kind.""" + def get_rules(self, exclude_scratchpad_when: bool = False) -> list[Engram]: + """Load behavioral rules (rules.md) as Engrams, grouped by kind. + + exclude_scratchpad_when drops "when" rules that mention "scratchpad" + — those are injected separately via recall_scratchpad_wisdom() into + the scratchpad tool description; including them in the system prompt + too would double their token cost. always/never rules are never + affected — those aren't part of that injection. + """ if not self._rules_path.is_file(): return [] try: @@ -390,6 +432,11 @@ def get_rules(self) -> list[Engram]: entries.sort(key=lambda x: x.kind) + if exclude_scratchpad_when: + entries = [ + e for e in entries if not (e.kind == "when" and _is_scratchpad_related(e)) + ] + return entries def del_rule(self, id): diff --git a/anton/core/runtime.py b/anton/core/runtime.py index 82783506..b078a898 100644 --- a/anton/core/runtime.py +++ b/anton/core/runtime.py @@ -141,32 +141,29 @@ async def build_chat_session( initial_history = history_store.load(session_id) data_vault = LocalDataVault() if LocalDataVault is not None else None - google_drive_oauth_connected = False if data_vault is not None: try: - for conn in data_vault.list_connections(): + connections = data_vault.list_connections() + except Exception: + logger.debug("Could not list Anton data vault connections", exc_info=True) + connections = [] + for conn in connections: + # Per-connection try/except so one bad record can't abort + # injecting env for the rest; engine/name are set before the try + # so the except's own logging can't itself raise. + engine = None + name = None + try: engine = conn.get("engine") name = conn.get("name") if not (engine and name): continue data_vault.inject_env(engine, name) - if engine == "google_drive": - fields = data_vault.load(engine, name) or {} - if fields.get("auth_type") == "oauth": - google_drive_oauth_connected = True - except Exception: - logger.debug("Could not inject Anton data vault env", exc_info=True) - - integration_guidance = "" - if google_drive_oauth_connected: - integration_guidance = ( - " Connected Google Drive accounts are available through Google OAuth credentials " - "in the injected `DS_GOOGLE_DRIVE___...` environment variables. " - "Only claim Google Drive access if you can actually use those credentials successfully." - ) - - suffix_parts = [s for s in (system_prompt_suffix, integration_guidance) if s] - final_suffix = "".join(suffix_parts) if suffix_parts else None + except Exception: + logger.debug( + "Could not inject Anton data vault env for %s/%s", + engine, name, exc_info=True, + ) config = ChatSessionConfig( llm_client=llm_client, @@ -176,7 +173,10 @@ async def build_chat_session( episodic=episodic, system_prompt_context=SystemPromptContext( runtime_context=build_runtime_context(settings), - suffix=final_suffix, + # SystemPromptContext.suffix is typed str (default ""), not + # Optional — pass "" rather than None when no suffix was given, + # or ChatSystemPromptBuilder.build()'s suffix.strip() crashes. + suffix=system_prompt_suffix or "", ), output_dir=str(output_dir), workspace=workspace, diff --git a/anton/core/session.py b/anton/core/session.py index 3f260aff..adbe78c8 100644 --- a/anton/core/session.py +++ b/anton/core/session.py @@ -2,13 +2,15 @@ import asyncio from collections.abc import AsyncIterator, Callable -from dataclasses import asdict, dataclass, field +from dataclasses import asdict, dataclass, field, replace from datetime import datetime import json import re -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING, List, Literal import os +from pydantic import BaseModel, Field + from anton.core.backends.base import Cell, ScratchpadRuntimeFactory from anton.core.backends.local import local_scratchpad_runtime_factory from anton.core.datasources.data_vault import DataVault @@ -120,6 +122,117 @@ def _scrub_user_input(user_input: str | list[dict]) -> str | list[dict]: ] +class _VerifierVerdict(BaseModel): + """Structured verdict from the completion verifier (runs on the cheap + coding model). The field descriptions below double as the verifier's + instructions — see LLMClient.generate_object_code (ENG-716).""" + + status: Literal["COMPLETE", "WAITING", "INCOMPLETE", "STUCK"] = Field( + description=( + "Classify the assistant's latest message against the user's request:\n" + "- COMPLETE: the requested task is done. A finished task followed by an " + "optional 'want me to…?' offer is still COMPLETE.\n" + "- WAITING: the assistant's latest message asks the user a question it " + "genuinely needs answered to proceed with the requested task, or is a " + "reasoned refusal. This is a valid stopping point — do NOT treat it as " + "unfinished; the correct action is to wait for the user's reply.\n" + "- INCOMPLETE: the assistant stopped partway through the requested task " + "WITHOUT asking the user anything, and could keep going on its own.\n" + "- STUCK: a hard blocker prevents completion (missing credentials, an " + "unavailable service, or a permission the assistant does not have)." + ) + ) + reason: str = Field(description="One brief sentence explaining the verdict.") + + +def _render_tool_result_content(content, cap: int) -> str: + """Render a tool_result's content as bounded plain text. + + Never serializes raw payloads: a multimodal result (e.g. read_image) can + carry megabytes of base64, so we keep only text blocks and mark images with + a placeholder rather than ``json.dumps``-ing the whole thing (ENG-716). + """ + if isinstance(content, str): + return content[:cap] or "(empty result)" + if isinstance(content, list): + parts: list[str] = [] + for block in content: + if not isinstance(block, dict): + continue + if block.get("type") == "text": + parts.append((block.get("text") or "").strip()) + elif block.get("type") in ("image", "image_url"): + parts.append("[image]") + return (" ".join(p for p in parts if p)[:cap]) or "[non-text result]" + return str(content)[:cap] + + +def _render_verify_transcript( + history: list[dict], + *, + max_convo: int = 10, + max_tool: int = 12, + tool_cap: int = 400, + text_cap: int = 2000, +) -> str: + """Render a compact, text-only view of the recent conversation for the + completion verifier. + + Budgets the conversational thread and the tool activity *separately* so a + voluminous tool loop can't crowd out the user/assistant turns the verifier + needs to resolve referential requests ("do the same for the other file"): + the most recent ``max_convo`` user/assistant text turns plus the most recent + ``max_tool`` tool events, merged back into chronological order. Speaker is + taken from the message role (list-based *user* content — images/files — is + labelled USER, not ASSISTANT); multimodal blocks are rendered block-by-block + (text kept, images as a placeholder, never raw base64); internal ``SYSTEM:`` + injections are dropped before budgeting so they don't consume slots. Keeps + the call cheap and free of tool_use/tool_result pairing constraints + (ENG-716). + """ + convo: list[tuple[int, str]] = [] # (orig_index, line) — user/assistant text + tools: list[tuple[int, str]] = [] # (orig_index, line) — tool activity + + for i, msg in enumerate(history): + role = msg.get("role") + content = msg.get("content") + speaker = "USER" if role == "user" else "ASSISTANT" + if isinstance(content, str): + text = content.strip() + if not text or (role == "user" and text.startswith("SYSTEM:")): + continue + convo.append((i, f"{speaker}: {text[:text_cap]}")) + elif isinstance(content, list): + # Assistant text emitted alongside a tool call is preamble/narration + # ("Processing step 4"), not a conversational turn — route it to the + # tool budget so a long tool loop can't evict real requests/replies + # from the conversation budget. + preamble = role == "assistant" and any( + isinstance(b, dict) and b.get("type") == "tool_use" for b in content + ) + for block in content: + if not isinstance(block, dict): + continue + btype = block.get("type") + if btype == "text": + text = (block.get("text") or "").strip() + if not text: + continue + bucket = tools if preamble else convo + bucket.append((i, f"{speaker}: {text[:text_cap]}")) + elif btype in ("image", "image_url"): + convo.append((i, f"{speaker}: [image]")) + elif btype == "tool_use": + tools.append((i, f"ASSISTANT called tool: {block.get('name')}")) + elif btype == "tool_result": + rendered = _render_tool_result_content(block.get("content"), tool_cap) + tools.append((i, f"TOOL RESULT: {rendered}")) + + kept = convo[-max_convo:] + tools[-max_tool:] + kept.sort(key=lambda entry: entry[0]) + return "\n".join(line for _, line in kept) or "(no conversation)" + + @dataclass class ChatSessionConfig: """All construction parameters for a ChatSession. @@ -183,6 +296,7 @@ def __init__(self, config: ChatSessionConfig) -> None: self._settings = config.settings self._max_tool_rounds = s.max_tool_rounds self._max_continuations = s.max_continuations + self._verify_min_tool_rounds = s.verify_min_tool_rounds self._context_pressure_threshold = s.context_pressure_threshold self._max_consecutive_errors = s.max_consecutive_errors self._resilience_nudge_at = s.resilience_nudge_at @@ -732,7 +846,10 @@ def _build_tools(self) -> list[dict]: return self.tool_registry.dump() def _build_core_tools(self) -> None: - scratchpad_tool = SCRATCHPAD_TOOL + # Copy — SCRATCHPAD_TOOL is a module-level singleton; mutating its + # .description in place would leak across every session/user sharing + # this process instead of resetting per session. + scratchpad_tool = replace(SCRATCHPAD_TOOL) pkg_list = self._scratchpads.available_packages if pkg_list: notable = sorted(p for p in pkg_list if p.lower() in self._NOTABLE_PACKAGES) @@ -1794,6 +1911,8 @@ async def _stream_and_handle_tools( # task isn't actually done yet. continuation = 0 _max_rounds_hit = False + import logging as _logging + _verifier_log = _logging.getLogger(__name__) while True: # Completion verification loop tool_round = 0 @@ -2194,9 +2313,10 @@ async def _stream_and_handle_tools( ) # --- Completion verification --- - # Only verify when tools were actually used (not for simple Q&A) - # and we haven't hit the max-rounds hard stop. - if tool_round == 0 or _max_rounds_hit: + # Skip when too few tool rounds were used (pure Q&A always skips at + # tool_round==0; raising verify_min_tool_rounds also skips trivial + # single-round turns) or when we hit the max-rounds hard stop. + if tool_round < self._verify_min_tool_rounds or _max_rounds_hit: break # Append the assistant's final text so the verifier can see it @@ -2227,47 +2347,64 @@ async def _stream_and_handle_tools( # Consolidation still runs after diagnosis break - # Ask the LLM to self-assess completion. - # Use a copy of history with a trailing user message so models - # that don't support assistant-prefill won't reject the request. - # Factory is re-invoked on each recovery attempt so the verifier - # sees the latest post-compaction history. - def build_verify_messages() -> list[dict]: - return list(self._history) + [ - { - "role": "user", - "content": ( - "SYSTEM: Evaluate whether the task the user originally requested " - "has been fully completed based on the conversation above." - ), - } - ] + # Ask the cheap coding model to self-assess completion over a compact, + # text-rendered view of the recent conversation: enough context for + # referential follow-ups plus truncated tool-result evidence to + # cross-check success claims, but far smaller than the raw transcript + # and free of tool_use/tool_result pairing constraints (ENG-716). The + # assistant's latest reply is already in history (appended above). + transcript = _render_verify_transcript(self._history) + # Always state the current request explicitly: a long tool-heavy turn + # can push the turn's opening user message out of the transcript window, + # and the request is the anchor for the whole judgment (ENG-716). + request = (user_message or "").strip() + request_header = f"USER'S CURRENT REQUEST: {request}\n\n" if request else "" + verify_messages = [ + { + "role": "user", + "content": ( + "Assess the conversation below (tool results are truncated) and " + "decide the status of the USER's most recent request.\n\n" + f"{request_header}{transcript}\n\n" + "Which status applies? A tool the task depended on that returned " + "an error — or returned no usable data while the assistant implied " + "success — means INCOMPLETE or STUCK, not COMPLETE." + ), + }, + ] verifier_system = ( - "You are a task-completion verifier. Given the conversation, determine " - "whether the user's original request has been fully completed.\n\n" - "Respond with EXACTLY one of these lines, followed by a brief reason:\n" - "STATUS: COMPLETE — \n" - "STATUS: INCOMPLETE — \n" - "STATUS: STUCK — \n\n" - "COMPLETE = the task is done or the response fully answers the question.\n" - "INCOMPLETE = more work can be done to finish the task.\n" - "STUCK = a blocker prevents completion (missing info, permissions, etc).\n\n" - "Be strict: if the user asked for X and only part of X was delivered, " - "that is INCOMPLETE, not COMPLETE. But if the user asked a question " - "and the assistant answered it, that is COMPLETE even without tool use." + "You are a task-completion verifier. Decide whether the user's " + "request is complete, the assistant is waiting on the user, the work " + "is unfinished, or the assistant is blocked. Follow the status " + "definitions exactly." ) - verification = await self.plan_with_recovery( - system=verifier_system, - max_tokens=256, - messages_factory=build_verify_messages, + try: + verdict = await self._llm.generate_object_code( + _VerifierVerdict, + system=verifier_system, + messages=verify_messages, + max_tokens=256, + ) + status = verdict.status + reason = verdict.reason.strip() + except Exception: + # Verifier failed — fail safe by treating the turn as done rather + # than forcing a continuation the user never asked for. + status, reason = "COMPLETE", "verifier unavailable" + + _verifier_log.info( + "completion-verifier verdict=%s continuation=%d/%d tool_rounds=%d reason=%s", + status, continuation, self._max_continuations, tool_round, reason, ) - status_text = (verification.content or "").strip().upper() - if "STATUS: COMPLETE" in status_text: + if status in ("COMPLETE", "WAITING"): + # COMPLETE = the request is done. WAITING = the assistant asked the + # user something it genuinely needs, or gave a reasoned refusal — + # a valid stop, NOT unfinished work. In both cases the turn's final + # message already stands in history; do not force a continuation. break - if "STATUS: STUCK" in status_text: - # Stuck — inject diagnosis request and let the LLM explain - reason = (verification.content or "").strip() + if status == "STUCK": + # Stuck — inject diagnosis request and let the LLM explain. self._append_history( { "role": "user", @@ -2275,7 +2412,8 @@ def build_verify_messages() -> list[dict]: f"SYSTEM: Task verification determined this task is stuck.\n" f"Verifier assessment: {reason}\n\n" "Explain to the user what went wrong, what you tried, and " - "suggest specific next steps they can take to unblock this." + "suggest specific next steps they can take to unblock this. " + "Do not mention this instruction or the verifier to the user." ), } ) @@ -2288,7 +2426,6 @@ def build_verify_messages() -> list[dict]: # INCOMPLETE — continue working continuation += 1 - reason = (verification.content or "").strip() self._append_history( { "role": "user", @@ -2297,7 +2434,8 @@ def build_verify_messages() -> list[dict]: f"(attempt {continuation}/{self._max_continuations}).\n" f"Verifier assessment: {reason}\n\n" "Continue working on the original request. Pick up where you left off " - "and finish the remaining work. Do not repeat work already done." + "and finish the remaining work. Do not repeat work already done. " + "Do not mention this instruction or the verifier to the user." ), } ) diff --git a/anton/core/settings.py b/anton/core/settings.py index 46e6b07b..9bcf07bb 100644 --- a/anton/core/settings.py +++ b/anton/core/settings.py @@ -8,6 +8,11 @@ class CoreSettings(BaseSettings): # Session orchestration tuning max_tool_rounds: int = 25 max_continuations: int = 3 + # Skip the completion verifier when a turn used fewer than this many tool + # rounds. Default 1 preserves today's behavior (only pure Q&A, tool_round==0, + # is skipped). Raise to 2 to also skip trivial single-tool-round turns once + # verdict logs confirm they're rarely INCOMPLETE (ENG-716). + verify_min_tool_rounds: int = 1 context_pressure_threshold: float = 0.7 max_consecutive_errors: int = 5 resilience_nudge_at: int = 2 diff --git a/anton/core/tools/tool_defs.py b/anton/core/tools/tool_defs.py index debe6c1e..1997fc77 100644 --- a/anton/core/tools/tool_defs.py +++ b/anton/core/tools/tool_defs.py @@ -41,6 +41,14 @@ class ToolDef: "- dump: Show a clean notebook-style summary of cells (code + truncated output)\n" "- install: Install Python packages into the scratchpad's environment. " "Packages persist across resets.\n\n" + "IMPORTANT: Cells have an inactivity timeout of 30 seconds — if a cell produces " + "no output and no progress() calls for 30s, it is killed and all state is lost. " + "For long-running code (API calls, data extraction, heavy computation), call " + "progress(message) periodically to signal work is ongoing and reset the timer. " + "The total timeout scales from your estimated_execution_time_seconds " + "(roughly 2x the estimate). You MUST provide estimated_execution_time_seconds " + "for every exec call. For very long operations, provide a realistic estimate " + "and use progress() to keep the cell alive.\n\n" "Use print() to produce output. Host Python packages are available by default. " "Include a 'packages' array on exec calls for any libraries your code needs — " "they'll be auto-installed before the cell runs (already-installed ones are skipped).\n" @@ -57,15 +65,7 @@ class ToolDef: "sample(var) inspects any variable with type-aware formatting — DataFrames get " "shape/dtypes/head, dicts get keys/values, lists get length/items. " "Defaults to 'preview' mode (compact); use sample(var, mode='full') for complete dump.\n" - "All .anton/.env secrets are available as environment variables (os.environ).\n\n" - "IMPORTANT: Cells have an inactivity timeout of 30 seconds — if a cell produces " - "no output and no progress() calls for 30s, it is killed and all state is lost. " - "For long-running code (API calls, data extraction, heavy computation), call " - "progress(message) periodically to signal work is ongoing and reset the timer. " - "The total timeout scales from your estimated_execution_time_seconds " - "(roughly 2x the estimate). You MUST provide estimated_execution_time_seconds " - "for every exec call. For very long operations, provide a realistic estimate " - "and use progress() to keep the cell alive." + "All .anton/.env secrets are available as environment variables (os.environ)." ), input_schema={ "type": "object", diff --git a/anton/utils/datasources.py b/anton/utils/datasources.py index 27b1f0a9..7354f035 100644 --- a/anton/utils/datasources.py +++ b/anton/utils/datasources.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import os import re import shutil @@ -7,7 +8,12 @@ from pathlib import Path from typing import TYPE_CHECKING -from anton.core.datasources.data_vault import DataVault, LocalDataVault, _slug_env_prefix +from anton.core.datasources.data_vault import ( + DataVault, + LocalDataVault, + _slug_env_prefix, + is_secret_key, +) from anton.core.datasources.datasource_registry import DatasourceRegistry, _YAML_BLOCK_RE if TYPE_CHECKING: @@ -104,6 +110,27 @@ def register_secret_vars( _DS_SECRET_VARS.add(key) +def _register_unregistered_connection_vars(vault: DataVault, engine: str, name: str) -> None: + """Register DS_* var names for a connection whose engine isn't in the + registry (custom engines, connector-spec saves). + + Without this, every field of such a connection falls into the + conservative unknown-DS_* scrub — so harmless values like base_url + surfaced as `[DS_..._BASE_URL]` markers in user-facing output (ENG-688). + Classification: the record's stored ``secure_keys`` when present, else + the vault's canonical legacy secret-name heuristic. + """ + record = vault.read_record(engine, name) if hasattr(vault, "read_record") else None + fields = (record or {}).get("fields") or vault.load(engine, name) or {} + secure_keys = (record or {}).get("secure_keys") + prefix = _slug_env_prefix(engine, name) + for field_name in fields: + key = f"{prefix}__{field_name.upper()}" + _DS_KNOWN_VARS.add(key) + if is_secret_key(field_name, secure_keys): + _DS_SECRET_VARS.add(key) + + def scrub_credentials(text: str) -> str: """Remove secret values from scratchpad/tool output before it reaches the LLM. @@ -145,6 +172,20 @@ def scrub_credentials(text: str) -> str: return text +def _parse_picked_files(raw: str | None) -> list[dict]: + """Parse a `_picked_files` JSON field, dropping malformed entries so they + can't crash a later f.get(...) call or embed "None" as a fileId.""" + if not raw: + return [] + try: + parsed = json.loads(raw) + except (json.JSONDecodeError, TypeError): + return [] + if not isinstance(parsed, list): + return [] + return [f for f in parsed if isinstance(f, dict) and f.get("id")] + + def build_datasource_context(vault: DataVault, active_only: str | None = None) -> str: """Build a system-prompt section listing available DS_* env vars by name. @@ -172,6 +213,12 @@ def build_datasource_context(vault: DataVault, active_only: str | None = None) - "it by name; never treat the bracket form as a literal credential " "or pass it back as a value to any tool.\n" ) + # Google Drive's drive.file OAuth scope only covers files the app created + # itself, plus files explicitly granted via the Google Picker (persisted + # as a `_picked_files` vault field) — these must be named explicitly here + # since a plain files.list()/files.search() call won't return them. + google_drive_oauth_connected = False + google_drive_picked_files: dict[str, list[dict]] = {} for c in conns: slug = f"{c['engine']}-{c['name']}" if active_only and slug != active_only: @@ -200,6 +247,39 @@ def build_datasource_context(vault: DataVault, active_only: str | None = None) - if identity: head += f" — {identity}" lines.append(f"- {head} → {var_names}") + if c["engine"] == "google_drive": + if fields.get("auth_type") == "oauth": + google_drive_oauth_connected = True + picked = _parse_picked_files(fields.get("_picked_files")) + if picked: + google_drive_picked_files[c["name"]] = picked + if google_drive_oauth_connected or google_drive_picked_files: + lines.append( + "\nConnected Google Drive accounts are available through Google OAuth credentials " + "in the injected `DS_GOOGLE_DRIVE___...` environment variables. " + "Only claim Google Drive access if you can actually use those credentials successfully." + ) + if google_drive_picked_files: + picked_lines = [ + f"- {f.get('name', 'untitled')} (id: {f.get('id')}, connection: {conn_name})" + for conn_name, files in google_drive_picked_files.items() + for f in files + ] + # Imperative and structurally separate from the paragraph above — a + # softer prose mention got silently dropped by the agent when + # reporting files.list() results verbatim. + lines.append( + "\nIMPORTANT — additional Drive files the user has explicitly granted access to " + "via the Google Picker, which a plain files.list() or files.search() call will NOT " + "return (the google_drive scope only covers files this app created itself, plus " + "these specifically granted ones):\n" + + "\n".join(picked_lines) + + "\nWhenever you list, search, or enumerate Drive files for the user, you MUST " + "include every file above IN ADDITION to whatever files.list()/files.search() " + "returns — do not report only the API call's results. To read one of these files' " + "content, call files.get(fileId=...) directly with its id above; do not expect it " + "to appear in a files.list() response first." + ) return "\n".join(lines) @@ -240,6 +320,8 @@ def restore_namespaced_env(vault: DataVault) -> None: edef = dreg.get(conn["engine"]) if edef is not None: register_secret_vars(edef, engine=conn["engine"], name=conn["name"]) + else: + _register_unregistered_connection_vars(vault, conn["engine"], conn["name"]) def find_matching_connection( diff --git a/docs/docs/developer/cerebellum-and-acc.md b/docs/docs/developer/cerebellum-and-acc.md index af8b647d..d7914e45 100644 --- a/docs/docs/developer/cerebellum-and-acc.md +++ b/docs/docs/developer/cerebellum-and-acc.md @@ -65,7 +65,9 @@ background task — the user gets their reply immediately: 3. Each lesson becomes an `Engram` with `kind="lesson"`, `topic="scratchpad"`, `source="consolidation"`, routed through `Cortex.encode()`. 4. Future scratchpad cells see those lessons via the existing - `recall_scratchpad_wisdom()` injection into the scratchpad tool description. + `recall_scratchpad_wisdom()` injection into the scratchpad tool description + (confidence tier + recency ordered, budget-limited — not all lessons are + guaranteed to survive the cut). The cerebellum is a **producer only** — no parallel storage system, no separate corrections file. Whatever the consolidator and `memorize` write to, the diff --git a/docs/docs/developer/memory-systems.md b/docs/docs/developer/memory-systems.md index 123c9553..552754da 100644 --- a/docs/docs/developer/memory-systems.md +++ b/docs/docs/developer/memory-systems.md @@ -56,7 +56,7 @@ it just reads and writes. | `recall_rules()` | `rules.md` | Basal Ganglia + OFC | | `recall_lessons(token_budget)` | `lessons.md` (budget-limited, most recent first) | Anterior Temporal Lobe | | `recall_topic(slug)` | `topics/{slug}.md` | Cortical Association Areas | -| `recall_scratchpad_wisdom()` | "when" rules + scratchpad-related lessons + `topics/scratchpad-*.md` | Procedural priming | +| `recall_scratchpad_wisdom(token_budget)` | scratchpad-related "when" rules + scratchpad-related lessons (confidence tier + recency ordered, budget-limited) | Procedural priming | | Encoding method | Writes | Behavior | |---|---|---| @@ -165,7 +165,9 @@ injected *after* memory, giving user instructions higher priority. **Moment B — scratchpad tool description (procedural priming).** When scratchpads are active, `cortex.get_scratchpad_context()` appends relevant lessons to the scratchpad tool's description — the LLM sees them exactly when -composing code. +composing code. Moment A excludes anything scratchpad-related +(`get_rules(exclude_scratchpad_when=True)`, `recall_lessons()`) so the same +rule/lesson isn't billed twice per call. ### Context budget diff --git a/tests/e2e/scenarios/test_loop_safety.py b/tests/e2e/scenarios/test_loop_safety.py index 25fdf5a6..3f83dcfe 100644 --- a/tests/e2e/scenarios/test_loop_safety.py +++ b/tests/e2e/scenarios/test_loop_safety.py @@ -50,6 +50,40 @@ def test_continuation_limit_respected(cfg, stub, tmp_path): ), f"Budget-exhausted message not found. Request count: {stub.request_count}" +@pytest.mark.stub_only +def test_waiting_verdict_stops_without_continuation(cfg, stub, tmp_path): + # ENG-716: a tool-using turn that ends by asking the user a question must + # STOP when the verifier returns WAITING — not inject "Continue working" + # and answer its own question. Also asserts the tool-outcome cross-check + # actually reaches the verifier. + stub.queue_tool_call("scratchpad", {"action": "exec", "name": "c", "code": "print(1)"}) + stub.queue_text("Which format would you like — PDF or HTML? WAITING_ON_USER") + stub.queue_verification_waiting("assistant asked the user a question it needs answered") + result = run_anton(["--folder", str(tmp_path)], ["make me a report", "exit"], + env=base_env(stub), timeout=cfg.timeout(30)) + + assert_exit_ok(result) + assert_not_output(result, "Traceback (most recent call last)") + assert_output(result, "WAITING_ON_USER") + # WAITING is a valid stop: no continuation injection may be sent. + assert not any( + "Continue working on the original request" in json.dumps(r.get("messages", [])) + for r in stub.requests + ), "WAITING verdict must not trigger a continuation injection" + # The verifier must receive truncated tool-result evidence (not just a flag), + # so it can cross-check claimed success against what the tool actually did. + assert any( + "TOOL RESULT:" in json.dumps(r.get("messages", [])) + for r in stub.requests + ), "verifier did not receive tool-result evidence" + # ...and the current request is always stated, even if a long turn evicts it + # from the transcript window. + assert any( + "USER'S CURRENT REQUEST" in json.dumps(r.get("messages", [])) + for r in stub.requests + ), "verifier did not receive the current request header" + + def test_session_exits_within_timeout(cfg, stub, tmp_path): stub.queue_text("Quick reply. QUICK_EXIT") stub.queue_verification_ok() diff --git a/tests/e2e/stub_server.py b/tests/e2e/stub_server.py index 6f812745..c8c20125 100644 --- a/tests/e2e/stub_server.py +++ b/tests/e2e/stub_server.py @@ -67,24 +67,35 @@ def queue_tool_call(self, name: str, arguments: dict) -> "StubServer": }])) return self - def queue_verification_ok(self) -> "StubServer": - """Queue a non-streaming 'STATUS: COMPLETE' verification response.""" - self._queue.put( - _Response(content="STATUS: COMPLETE — task is done.", force_streaming=False) - ) + def _queue_verdict(self, status: str, reason: str) -> "StubServer": + """Queue a non-streaming structured verifier verdict. + + The completion verifier now runs via ``generate_object_code`` with a + forced tool_choice on the ``_VerifierVerdict`` schema, so the stub must + answer with a tool call (status + reason), not free text (ENG-716). + """ + self._queue.put(_Response( + tool_calls=[{ + "id": f"call_{uuid.uuid4().hex[:8]}", + "name": "_VerifierVerdict", + "arguments": {"status": status, "reason": reason}, + }], + force_streaming=False, + )) return self + def queue_verification_ok(self) -> "StubServer": + """Queue a COMPLETE verifier verdict.""" + return self._queue_verdict("COMPLETE", "task is done.") + def queue_verification_incomplete(self, reason: str = "still more to do") -> "StubServer": - self._queue.put( - _Response(content=f"STATUS: INCOMPLETE — {reason}", force_streaming=False) - ) - return self + return self._queue_verdict("INCOMPLETE", reason) + + def queue_verification_waiting(self, reason: str = "waiting on the user") -> "StubServer": + return self._queue_verdict("WAITING", reason) def queue_verification_stuck(self, reason: str = "blocked") -> "StubServer": - self._queue.put( - _Response(content=f"STATUS: STUCK — {reason}", force_streaming=False) - ) - return self + return self._queue_verdict("STUCK", reason) def queue_summary(self, text: str = "Summary of earlier turns.") -> "StubServer": """Queue a response for _summarize_history's coding model call.""" @@ -212,6 +223,16 @@ def _send_sse(handler: BaseHTTPRequestHandler, resp: _Response) -> None: def _send_json(handler: BaseHTTPRequestHandler, resp: _Response) -> None: + message: dict = {"role": "assistant", "content": resp.content or None} + finish_reason = "stop" + if resp.tool_calls: + tc = resp.tool_calls[0] + message["tool_calls"] = [{ + "id": tc["id"], + "type": "function", + "function": {"name": tc["name"], "arguments": json.dumps(tc["arguments"])}, + }] + finish_reason = "tool_calls" data = { "id": f"chatcmpl-{uuid.uuid4().hex[:8]}", "object": "chat.completion", @@ -219,8 +240,8 @@ def _send_json(handler: BaseHTTPRequestHandler, resp: _Response) -> None: "model": "gpt-test", "choices": [{ "index": 0, - "message": {"role": "assistant", "content": resp.content}, - "finish_reason": "stop", + "message": message, + "finish_reason": finish_reason, }], "usage": {"prompt_tokens": 10, "completion_tokens": 10, "total_tokens": 20}, } diff --git a/tests/test_build_chat_session_google_drive.py b/tests/test_build_chat_session_google_drive.py new file mode 100644 index 00000000..20e86b29 --- /dev/null +++ b/tests/test_build_chat_session_google_drive.py @@ -0,0 +1,84 @@ +"""End-to-end: build_chat_session() -> ChatSession._build_system_prompt() actually +surfaces Google Drive Picker guidance in the real assembled system prompt. + +ENG-687 review (PR #241): the picked-files/OAuth guidance moved out of +build_chat_session (anton/core/runtime.py) into build_datasource_context +(anton/utils/datasources.py), which ChatSession already calls fresh every +turn. This exercises the real integration point end-to-end rather than just +the relocated function in isolation, and guards a real crash the move +uncovered: SystemPromptContext.suffix is typed str (default ""), but +build_chat_session could pass suffix=None when no system_prompt_suffix was +given, and ChatSystemPromptBuilder.build() calls suffix.strip() unconditionally. +""" +from __future__ import annotations + +import json + +import pytest + + +@pytest.fixture(autouse=True) +def _isolated_home(tmp_path, monkeypatch): + """LocalDataVault() with no args resolves under $HOME — isolate it so this + test never touches the real ~/.anton/data_vault or ~/.anton/memory.""" + home = tmp_path / "home" + home.mkdir() + monkeypatch.setenv("HOME", str(home)) + # build_chat_session constructs a real LLMClient; a syntactically-plausible + # dummy key is enough since this test never calls the LLM (no turn_stream()). + monkeypatch.setenv("ANTON_ANTHROPIC_API_KEY", "sk-ant-test-00000000000000000000000000") + return home + + +@pytest.fixture() +def workspace_path(tmp_path): + p = tmp_path / "workspace" + p.mkdir() + return p + + +async def test_picked_files_reach_the_real_system_prompt(workspace_path): + from anton.core.datasources.data_vault import LocalDataVault + from anton.core.runtime import build_chat_session + + vault = LocalDataVault() + vault.save("google_drive", "work", { + "auth_type": "oauth", + "account_email": "user@example.com", + "_picked_files": json.dumps([{"id": "f1", "name": "Roadmap.gdoc"}]), + }) + + session = await build_chat_session(session_id="test-picked-files", workspace_path=str(workspace_path)) + prompt = await session._build_system_prompt() + + assert "Connected Google Drive accounts are available" in prompt + assert "IMPORTANT — additional Drive files" in prompt + assert "Roadmap.gdoc" in prompt + + +async def test_no_connections_and_no_suffix_does_not_crash(workspace_path): + """Regression: on staging, build_chat_session could pass suffix=None to + SystemPromptContext (typed str) whenever there was nothing to add, + crashing ChatSystemPromptBuilder.build()'s suffix.strip() call. Any + session with zero connections and no explicit system_prompt_suffix hits + this — verifying it explicitly rather than relying on it being masked by + an unrelated google_drive connection existing in the vault.""" + from anton.core.runtime import build_chat_session + + session = await build_chat_session(session_id="test-empty", workspace_path=str(workspace_path)) + prompt = await session._build_system_prompt() # must not raise + + assert "Google Drive" not in prompt + + +async def test_explicit_system_prompt_suffix_still_appended(workspace_path): + from anton.core.runtime import build_chat_session + + session = await build_chat_session( + session_id="test-suffix", + workspace_path=str(workspace_path), + system_prompt_suffix="Host-specific note: reply in French.", + ) + prompt = await session._build_system_prompt() + + assert "Host-specific note: reply in French." in prompt diff --git a/tests/test_chat_scratchpad.py b/tests/test_chat_scratchpad.py index e7175f59..600b056e 100644 --- a/tests/test_chat_scratchpad.py +++ b/tests/test_chat_scratchpad.py @@ -7,7 +7,7 @@ import pytest from anton.core.backends.base import Cell -from anton.core.session import ChatSession, ChatSessionConfig +from anton.core.session import ChatSession, ChatSessionConfig, _VerifierVerdict from anton.core.tools.tool_defs import SCRATCHPAD_TOOL from anton.commands.session import handle_resume from anton.core.llm.provider import LLMResponse, StreamComplete, StreamToolResult, ToolCall, Usage @@ -85,6 +85,38 @@ async def test_scratchpad_tool_in_tools(self, workspace): finally: await session.close() + async def test_tool_build_does_not_mutate_shared_singleton(self, workspace): + """_build_core_tools() must not mutate the shared SCRATCHPAD_TOOL + singleton — otherwise memory-injected wisdom would leak across every + session sharing this process instead of resetting per session.""" + original_description = SCRATCHPAD_TOOL.description + + mock_llm = make_mock_llm() + cortex_a = MagicMock() + cortex_a.get_scratchpad_context.return_value = "WISDOM_FROM_SESSION_A" + session_a = ChatSession( + ChatSessionConfig(llm_client=mock_llm, workspace=workspace, cortex=cortex_a) + ) + try: + tools_a = session_a._build_tools() + desc_a = next(t["description"] for t in tools_a if t["name"] == "scratchpad") + assert "WISDOM_FROM_SESSION_A" in desc_a + assert SCRATCHPAD_TOOL.description == original_description + finally: + await session_a.close() + + cortex_b = MagicMock() + cortex_b.get_scratchpad_context.return_value = "" + session_b = ChatSession( + ChatSessionConfig(llm_client=mock_llm, workspace=workspace, cortex=cortex_b) + ) + try: + tools_b = session_b._build_tools() + desc_b = next(t["description"] for t in tools_b if t["name"] == "scratchpad") + assert "WISDOM_FROM_SESSION_A" not in desc_b + finally: + await session_b.close() + class TestScratchpadExecViaChat: async def test_scratchpad_exec_via_chat(self, workspace): @@ -219,7 +251,9 @@ async def test_scratchpad_dump_streams_tool_result(self, workspace): """dump action yields a StreamToolResult for display, but sends a short summary back to the LLM to avoid it parroting the full notebook.""" mock_llm = make_mock_llm() - mock_llm.plan = AsyncMock(return_value=_text_response("STATUS: COMPLETE — task done")) + mock_llm.generate_object_code = AsyncMock( + return_value=_VerifierVerdict(status="COMPLETE", reason="task done") + ) call_count = 0 @@ -274,7 +308,9 @@ async def test_scratchpad_in_streaming_path(self, workspace): final_response = _text_response("Got 99.") mock_llm = make_mock_llm() - mock_llm.plan = AsyncMock(return_value=_text_response("STATUS: COMPLETE — task done")) + mock_llm.generate_object_code = AsyncMock( + return_value=_VerifierVerdict(status="COMPLETE", reason="task done") + ) call_count = 0 diff --git a/tests/test_cortex.py b/tests/test_cortex.py index 69facbec..28faa86d 100644 --- a/tests/test_cortex.py +++ b/tests/test_cortex.py @@ -59,6 +59,35 @@ async def test_includes_lessons(self, cortex, dirs): assert "Global fact" in result assert "Project fact" in result + async def test_excludes_scratchpad_when_rules(self, cortex, dirs): + """Scratchpad-related "when" rules already surface via the scratchpad + tool description (get_scratchpad_context) — showing them here too + would double their token cost, so they're excluded from the system + prompt. Unrelated "when" rules are unaffected.""" + g, _ = dirs + hc = Hippocampus(g) + hc.encode_rule( + "If a scratchpad API is paginated, use progress()", + kind="when", confidence="high", source="user", + ) + hc.encode_rule( + "If the user writes in Spanish, respond in Spanish", + kind="when", confidence="high", source="user", + ) + result = await cortex.build_memory_context() + assert "paginated" not in result + assert "Spanish" in result + + async def test_excludes_scratchpad_lessons(self, cortex, dirs): + """Same exclusion as above, for lessons.""" + g, _ = dirs + hc = Hippocampus(g) + hc.encode_lesson("Scratchpad cells timeout at 30s") + hc.encode_lesson("CoinGecko rate-limits at 50/min") + result = await cortex.build_memory_context() + assert "timeout at 30s" not in result + assert "rate-limits" in result + class TestGetScratchpadContext: def test_empty_returns_empty(self, cortex): @@ -66,7 +95,9 @@ def test_empty_returns_empty(self, cortex): def test_combines_scopes(self, cortex, dirs): g, p = dirs - (g / "rules.md").write_text("# Rules\n\n## Always\n\n## Never\n\n## When\n- If slow → batch\n") + (g / "rules.md").write_text( + "# Rules\n\n## Always\n\n## Never\n\n## When\n- If scratchpad is slow → batch\n" + ) (p / "lessons.md").write_text("# Lessons\n- Scratchpad times out at 30s\n") result = cortex.get_scratchpad_context() assert "slow" in result diff --git a/tests/test_datasource_context_identity.py b/tests/test_datasource_context_identity.py index dff4d778..667ea0e2 100644 --- a/tests/test_datasource_context_identity.py +++ b/tests/test_datasource_context_identity.py @@ -5,8 +5,10 @@ or the DB host/name, never the credential, and never a dump of opaque/config fields. """ +import json + from anton.core.datasources.data_vault import LocalDataVault -from anton.utils.datasources import _connection_identity, build_datasource_context +from anton.utils.datasources import _connection_identity, _parse_picked_files, build_datasource_context class TestConnectionIdentity: @@ -82,3 +84,106 @@ def test_label_preferred_over_email(self, tmp_path): ctx = build_datasource_context(v) assert "Support" in ctx # the human label is shown assert "regtr@mail.com" not in ctx # label preferred over the opaque email + + +class TestParsePickedFiles: + def test_valid_list(self): + raw = json.dumps([{"id": "1", "name": "a"}, {"id": "2", "name": "b"}]) + assert _parse_picked_files(raw) == [{"id": "1", "name": "a"}, {"id": "2", "name": "b"}] + + def test_drops_malformed_entries(self): + raw = json.dumps([ + {"id": "1", "name": "keep"}, + "not-a-dict", + {"name": "missing-id"}, + {"id": None, "name": "null-id"}, + ]) + assert _parse_picked_files(raw) == [{"id": "1", "name": "keep"}] + + def test_non_list_json_returns_empty(self): + assert _parse_picked_files(json.dumps({"id": "1"})) == [] + + def test_invalid_json_returns_empty(self): + assert _parse_picked_files("not json") == [] + + def test_empty_or_none_returns_empty(self): + assert _parse_picked_files(None) == [] + assert _parse_picked_files("") == [] + + +class TestGoogleDrivePickerContext: + """ENG-687: google_drive's drive.file OAuth scope only covers files the + app created itself, plus files explicitly granted via the Google + Picker — the agent needs those named by id or a plain files.list()/ + files.search() call won't surface them at all.""" + + def test_oauth_connection_without_picked_files_shows_availability_only(self, tmp_path): + v = LocalDataVault(tmp_path) + v.save("google_drive", "work", {"auth_type": "oauth", "account_email": "u@x.com"}) + ctx = build_datasource_context(v) + assert "Connected Google Drive accounts are available" in ctx + assert "IMPORTANT" not in ctx + + def test_picked_files_surfaced_with_id_and_connection(self, tmp_path): + v = LocalDataVault(tmp_path) + v.save("google_drive", "work", { + "auth_type": "oauth", + "_picked_files": json.dumps([{"id": "f1", "name": "Roadmap.gdoc"}]), + }) + ctx = build_datasource_context(v) + assert "IMPORTANT — additional Drive files" in ctx + assert "Roadmap.gdoc" in ctx + assert "id: f1" in ctx + assert "connection: work" in ctx + + def test_resource_key_not_required_but_included_when_present(self, tmp_path): + v = LocalDataVault(tmp_path) + v.save("google_drive", "work", { + "_picked_files": json.dumps([{"id": "f1", "name": "Shared.gdoc", "resourceKey": "rk123"}]), + }) + ctx = build_datasource_context(v) + assert "Roadmap.gdoc" not in ctx # sanity: not leaking the other test's fixture + assert "Shared.gdoc" in ctx + + def test_malformed_picked_file_entries_are_dropped_not_crashed(self, tmp_path): + v = LocalDataVault(tmp_path) + v.save("google_drive", "work", { + "_picked_files": json.dumps(["not-a-dict", {"name": "missing-id"}]), + }) + ctx = build_datasource_context(v) # must not raise + assert "IMPORTANT" not in ctx # nothing well-formed survived, so no block at all + + def test_no_google_drive_connection_no_guidance(self, tmp_path): + v = LocalDataVault(tmp_path) + v.save("postgres", "prod", {"host": "db.acme.com"}) + ctx = build_datasource_context(v) + assert "Google Drive" not in ctx + assert "IMPORTANT" not in ctx + + def test_multiple_google_drive_connections_each_listed_separately(self, tmp_path): + v = LocalDataVault(tmp_path) + v.save("google_drive", "work", {"_picked_files": json.dumps([{"id": "1", "name": "Work.gdoc"}])}) + v.save("google_drive", "personal", {"_picked_files": json.dumps([{"id": "2", "name": "Personal.gdoc"}])}) + ctx = build_datasource_context(v) + assert "Work.gdoc" in ctx and "connection: work" in ctx + assert "Personal.gdoc" in ctx and "connection: personal" in ctx + + def test_active_only_suppresses_other_connections_guidance(self, tmp_path): + v = LocalDataVault(tmp_path) + v.save("google_drive", "work", { + "auth_type": "oauth", + "_picked_files": json.dumps([{"id": "1", "name": "Roadmap.gdoc"}]), + }) + v.save("postgres", "prod", {"host": "db.acme.com"}) + # A different connection is active — Drive guidance must not leak in. + ctx = build_datasource_context(v, active_only="postgres-prod") + assert "Google Drive" not in ctx + assert "Roadmap.gdoc" not in ctx + + def test_active_only_on_google_drive_still_shows_its_guidance(self, tmp_path): + v = LocalDataVault(tmp_path) + v.save("google_drive", "work", { + "_picked_files": json.dumps([{"id": "1", "name": "Roadmap.gdoc"}]), + }) + ctx = build_datasource_context(v, active_only="google_drive-work") + assert "Roadmap.gdoc" in ctx diff --git a/tests/test_hippocampus.py b/tests/test_hippocampus.py index 2d3036db..712433d0 100644 --- a/tests/test_hippocampus.py +++ b/tests/test_hippocampus.py @@ -80,10 +80,13 @@ def test_empty_returns_empty(self, hc): def test_extracts_when_rules(self, hc, mem_dir): (mem_dir / "rules.md").write_text( - "# Rules\n\n## Always\n- Be fast\n\n## When\n- If paginated → use progress()\n" + "# Rules\n\n## Always\n- Be fast\n\n## When\n" + "- If a scratchpad API is paginated → use progress()\n" + "- If the user writes in Spanish → respond in Spanish\n" ) result = hc.recall_scratchpad_wisdom() assert "paginated" in result + assert "Spanish" not in result def test_includes_scratchpad_lessons(self, hc, mem_dir): (mem_dir / "lessons.md").write_text( @@ -98,6 +101,37 @@ def test_includes_scratchpad_topic_files(self, hc, mem_dir): result = hc.recall_scratchpad_wisdom() assert "Always re-import" in result + def test_sorts_by_confidence_tier_then_recency(self, hc, mem_dir): + (mem_dir / "rules.md").write_text( + "# Rules\n\n## When\n" + "- Scratchpad HIGH_OLD_RULE \n" + "- Scratchpad LOW_NEW_RULE \n" + "- Scratchpad MEDIUM_MID_RULE \n" + ) + (mem_dir / "lessons.md").write_text( + "# Lessons\n" + "- Scratchpad HIGH_NEW_LESSON fact \n" + ) + result = hc.recall_scratchpad_wisdom() + + # High tier (newest first), then medium, then low last. + assert ( + result.index("HIGH_NEW_LESSON") + < result.index("HIGH_OLD_RULE") + < result.index("MEDIUM_MID_RULE") + < result.index("LOW_NEW_RULE") + ) + + def test_budget_limits_output(self, hc, mem_dir): + entries = "\n".join( + f"- Scratchpad when-rule number {i} with some extra padding words here" + for i in range(30) + ) + (mem_dir / "rules.md").write_text(f"# Rules\n\n## When\n{entries}\n") + result = hc.recall_scratchpad_wisdom(token_budget=150) + entry_count = result.count("- Scratchpad when-rule") + assert 0 < entry_count < 30 + class TestEncodeRule: def test_creates_rules_file(self, hc, mem_dir): diff --git a/tests/test_scrubbing.py b/tests/test_scrubbing.py index a7cc4772..87518a90 100644 --- a/tests/test_scrubbing.py +++ b/tests/test_scrubbing.py @@ -139,3 +139,71 @@ def test_known_secret_env_value_redacted_with_label(self, monkeypatch): result = _scrub_user_input(f"my key is {key}") assert key not in result assert "[OPENAI_API_KEY]" in result + + +class TestCustomEngineRegistration: + """ENG-688: connections of engines not in the registry (custom engines, + connector-spec saves) must register their fields so non-secret values + (base_url, host, ...) stay readable instead of leaking as markers.""" + + def _vault(self, tmp_path): + from anton.core.datasources.data_vault import LocalDataVault + + return LocalDataVault(tmp_path / "vault") + + def test_custom_engine_base_url_readable_secret_scrubbed(self, tmp_path): + from anton.utils.datasources import restore_namespaced_env + + vault = self._vault(tmp_path) + vault.save( + "acme_crm", "prod", + {"base_url": "https://api.acme-crm.example", "token": "tok_1234567890abcdef"}, + secure_keys=["token"], + ) + restore_namespaced_env(vault) + + result = scrub_credentials( + "GET https://api.acme-crm.example failed with token tok_1234567890abcdef" + ) + assert "https://api.acme-crm.example" in result + assert "tok_1234567890abcdef" not in result + assert "[DS_ACME_CRM_PROD__TOKEN]" in result + + def test_custom_engine_without_secure_keys_uses_name_heuristic(self, tmp_path): + from anton.utils.datasources import restore_namespaced_env + + vault = self._vault(tmp_path) + vault.save( + "acme_crm", "legacy", + {"base_url": "https://legacy.acme-crm.example", "api_key": "ak_1234567890abcdef"}, + ) + restore_namespaced_env(vault) + + result = scrub_credentials( + "base https://legacy.acme-crm.example key ak_1234567890abcdef" + ) + assert "https://legacy.acme-crm.example" in result + assert "ak_1234567890abcdef" not in result + assert "[DS_ACME_CRM_LEGACY__API_KEY]" in result + + def test_custom_engine_legacy_passphrase_is_scrubbed(self, tmp_path): + from anton.utils.datasources import restore_namespaced_env + + vault = self._vault(tmp_path) + passphrase = "correct horse battery staple" + vault.save( + "acme_crm", + "legacy", + { + "base_url": "https://legacy.acme-crm.example", + "passphrase": passphrase, + }, + ) + restore_namespaced_env(vault) + + result = scrub_credentials( + f"base https://legacy.acme-crm.example passphrase {passphrase}" + ) + assert "https://legacy.acme-crm.example" in result + assert passphrase not in result + assert "[DS_ACME_CRM_LEGACY__PASSPHRASE]" in result diff --git a/tests/test_verify_transcript.py b/tests/test_verify_transcript.py new file mode 100644 index 00000000..1d7ba6a9 --- /dev/null +++ b/tests/test_verify_transcript.py @@ -0,0 +1,134 @@ +"""Unit tests for the completion-verifier's compact transcript view (ENG-716). + +Guards the two review findings on PR #255: the verifier view must retain +referential task context from prior turns AND carry truncated tool-result +evidence (not just an ok/error flag), while omitting internal SYSTEM +injections. +""" + +from __future__ import annotations + +from anton.core.session import _render_verify_transcript + + +def test_keeps_prior_turn_context_for_referential_followups(): + history = [ + {"role": "user", "content": "clean up report_q1.csv"}, + {"role": "assistant", "content": "Done — deduped 12 rows."}, + {"role": "user", "content": "now do the same for the other file"}, + {"role": "assistant", "content": "Which file did you mean?"}, + ] + out = _render_verify_transcript(history) + # The current request is referential ("the same", "the other file"); the + # prior turn must remain so the verifier can resolve it. + assert "clean up report_q1.csv" in out + assert "now do the same for the other file" in out + + +def test_includes_truncated_tool_result_evidence(): + history = [ + {"role": "user", "content": "how many open PRs?"}, + {"role": "assistant", "content": [{"type": "tool_use", "name": "scratchpad"}]}, + {"role": "user", "content": [{"type": "tool_result", "content": "[error]\nHTTP 403 Forbidden"}]}, + {"role": "assistant", "content": "There are 42 open PRs."}, + ] + out = _render_verify_transcript(history) + # The tool actually errored while the assistant claimed a number — the + # verifier must see the error content to catch the mismatch. + assert "TOOL RESULT: [error]" in out + assert "HTTP 403 Forbidden" in out + assert "There are 42 open PRs." in out + + +def test_omits_internal_system_injections(): + history = [ + {"role": "user", "content": "build the dashboard"}, + {"role": "user", "content": "SYSTEM: Task verification determined this task is not yet complete. Continue working."}, + {"role": "assistant", "content": "Working on it."}, + ] + out = _render_verify_transcript(history) + assert "SYSTEM:" not in out + assert "Continue working" not in out + assert "build the dashboard" in out + + +def test_truncates_large_tool_output(): + big = "x" * 5000 + history = [ + {"role": "user", "content": "run it"}, + {"role": "user", "content": [{"type": "tool_result", "content": big}]}, + ] + out = _render_verify_transcript(history, tool_cap=400) + # Tool output is capped so the verify call stays cheap. + assert out.count("x") <= 400 + + +def test_empty_history_is_safe(): + assert _render_verify_transcript([]) == "(no conversation)" + + +def test_long_tool_turn_keeps_request_and_antecedent(): + # A prior turn establishes context; the current turn is referential and runs + # >8 tool rounds. Tool activity must NOT evict the conversational thread. + history = [ + {"role": "user", "content": "clean up report_q1.csv"}, + {"role": "assistant", "content": "Done — deduped 12 rows in report_q1.csv."}, + {"role": "user", "content": "now do the same for the other file"}, + ] + for i in range(10): # 10 tool rounds, each with preamble text + a tool call + history.append({"role": "assistant", "content": [ + {"type": "text", "text": f"Processing step {i}"}, + {"type": "tool_use", "name": "scratchpad"}, + ]}) + history.append({"role": "user", "content": [{"type": "tool_result", "content": f"row {i} cleaned"}]}) + history.append({"role": "assistant", "content": "All cleaned."}) + + out = _render_verify_transcript(history) + # Both the referential request and its antecedent survive the tool volume — + # preamble text ("Processing step N") must not consume the conversation budget. + assert "now do the same for the other file" in out + assert "clean up report_q1.csv" in out + assert "Done — deduped 12 rows in report_q1.csv." in out # antecedent's answer + # And tool evidence is still present. + assert "TOOL RESULT:" in out + + +def test_multimodal_user_text_labeled_user_not_assistant(): + history = [ + {"role": "user", "content": [ + {"type": "text", "text": "describe this image and save a report"}, + {"type": "image", "source": {"type": "base64", "data": "A" * 500}}, + ]}, + {"role": "assistant", "content": "Here's the description."}, + ] + out = _render_verify_transcript(history) + assert "USER: describe this image and save a report" in out + assert "ASSISTANT: describe this image" not in out + assert "USER: [image]" in out + assert "A" * 100 not in out # no base64 leaked + + +def test_multimodal_tool_result_keeps_summary_drops_base64(): + history = [ + {"role": "user", "content": "read the chart"}, + {"role": "user", "content": [{"type": "tool_result", "content": [ + {"type": "image", "source": {"type": "base64", "data": "Z" * 5000}}, + {"type": "text", "text": "Chart saved: 12 monthly bars"}, + ]}]}, + ] + out = _render_verify_transcript(history) + assert "Chart saved: 12 monthly bars" in out + assert "[image]" in out + assert "Z" * 100 not in out # base64 payload never serialized into the view + + +def test_system_injections_do_not_consume_budget(): + # SYSTEM injections are dropped before budgeting, so real turns aren't evicted + # by internal noise. + history = [] + for i in range(8): + history.append({"role": "user", "content": f"SYSTEM: internal note {i}"}) + history.append({"role": "user", "content": "the actual request"}) + out = _render_verify_transcript(history, max_convo=3) + assert "the actual request" in out + assert "internal note" not in out diff --git a/uv.lock b/uv.lock index 08c2abc4..e23502f3 100644 --- a/uv.lock +++ b/uv.lock @@ -203,7 +203,7 @@ requires-dist = [ { name = "pytest", marker = "extra == 'dev'", specifier = ">=9.0.0" }, { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.24" }, { name = "pyyaml", specifier = ">=6.0" }, - { name = "rich", specifier = ">=13.0" }, + { name = "rich", specifier = ">=13.0,<15" }, { name = "typer", specifier = ">=0.12.4" }, ] provides-extras = ["clipboard", "dev"]