From 71aa2e6224a26852e533d769c386f2fcd951f35e Mon Sep 17 00:00:00 2001 From: Emanuilo Jovanovic Date: Wed, 7 Oct 2026 21:26:31 +0200 Subject: [PATCH 01/10] feat: persistent Claude sessions, /new, /model, /effort - Reuse one Claude CLI client per chat instead of spawning one per turn - Resume conversations after restart (sessions.json), fresh session on failure - Replace the rolling history window with native sessions; /new resets - Static system prompt; time and memory context move into the user message - Load no ~/.claude settings/plugins/MCP servers, no built-in tools, dontAsk - Telegram /model and /effort, persisted to ~/.memclaw/.env - Background consolidation after the reply - Default to claude-sonnet-5-5 at medium effort; bump claude-agent-sdk --- .env.example | 2 + README.md | 20 +- docs/docker.md | 15 +- docs/memory-architecture.md | 90 +++--- memclaw/agent.py | 318 +++++++++++++------- memclaw/backends/base.py | 64 +++- memclaw/backends/claude.py | 387 +++++++++++++++++++----- memclaw/backends/cursor.py | 64 +++- memclaw/bot/handlers.py | 111 ++++++- memclaw/bot/slack_handlers.py | 7 +- memclaw/bot/whatsapp_handlers.py | 7 +- memclaw/cli.py | 31 +- memclaw/config.py | 4 - memclaw/defaults/AGENTS.md | 2 +- memclaw/sessions.py | 65 ++++ memclaw/setup.py | 23 +- pyproject.toml | 2 +- tests/backends/test_claude.py | 453 ++++++++++++++++++++++++---- tests/backends/test_cursor.py | 62 ++++ tests/test_agent.py | 493 ++++++++++++++++++++----------- tests/test_config.py | 28 +- tests/test_handlers.py | 200 +++++++++++++ uv.lock | 16 +- 23 files changed, 1957 insertions(+), 507 deletions(-) create mode 100644 memclaw/sessions.py diff --git a/.env.example b/.env.example index d6fd2de..49f7461 100644 --- a/.env.example +++ b/.env.example @@ -7,6 +7,8 @@ OPENAI_API_KEY=your-openai-key # Required for the Claude agent (interactive mode + /ask command) ANTHROPIC_API_KEY=your-anthropic-key # AGENT_BACKEND=claude +# CLAUDE_MODEL=claude-sonnet-5-5 +# CLAUDE_EFFORT=medium # Required for Cursor agent # AGENT_BACKEND=cursor diff --git a/README.md b/README.md index fb486f8..0b429d1 100644 --- a/README.md +++ b/README.md @@ -70,10 +70,14 @@ The main way to use Memclaw. Just talk to it naturally — no commands needed. S All platforms share the same agent, memories, and search index — your data is unified regardless of how you interact. +Each chat is one ongoing conversation that survives restarts. Send `/new` to start a fresh one (your memories are kept). On Telegram it's a menu command; in Slack, WhatsApp and the terminal send a message that is exactly `/new`. + ### Telegram Bot The Telegram bot shows a **typing indicator** while processing so you know it's working on your request. +Commands: `/new` starts a fresh conversation, `/model` lists the available models (`/model ` switches), and `/effort` shows or sets the reasoning effort. Model and effort changes are saved to `~/.memclaw/.env` and keep the current conversation. They're available with the Claude backend. + #### Setup 1. Create a bot via [@BotFather](https://t.me/BotFather) and copy the token. @@ -208,7 +212,9 @@ The wizard asks you to pick one of two auth modes: Only one of the two is set at a time. -The wizard then lets you pick a model from the list Anthropic offers your credential and, if the model supports it, an effort level. Pressing Enter keeps your current choice. +The wizard then lets you pick a model from the list Anthropic offers your credential and, if the model supports it, an effort level. Pressing Enter keeps your current choice. The default is `claude-sonnet-5-5` at `medium` effort. + +Each chat keeps one Claude Code process running and reuses it, so replies don't pay a start-up cost, and the conversation lives in a Claude Code session that is resumed after a restart. The CLI runs without your `~/.claude` settings, plugins or MCP servers, and with no built-in tools — only Memclaw's own. ### Cursor @@ -259,7 +265,7 @@ Every memory is chunked, embedded, and indexed in SQLite. Retrieval combines two ### Agent Layer -A pluggable agent SDK drives the loop — Claude or Cursor, picked during setup ([Agent Backend](#agent-backend)). The agent keeps a rolling conversation history that always covers at least the last 10 message pairs, plus any older messages from the last hour — whichever set is larger — so short bursts of chatter stay in context without dropping continuity during slower conversations. It decides when to **store** vs **search** based on your intent. +A pluggable agent SDK drives the loop — Claude or Cursor, picked during setup ([Agent Backend](#agent-backend)). Each chat's conversation lives in the backend's own session, resumed after a restart; the system prompt (AGENTS.md + MEMORY.md) stays fixed per session so it can be prompt-cached, and per-message context — the current time and relevant memories — rides at the top of each message. It decides when to **store** vs **search** based on your intent. | Tool | What it does | |------|-------------| @@ -289,6 +295,9 @@ Chat naturally: Based on your memories, Alex is someone you had coffee with recently. She's moving to Berlin for a new role at Stripe. +> /new +Started a new conversation. + > /quit ``` @@ -348,8 +357,8 @@ Backend choice and credentials are covered above in [Agent Backend](#agent-backe | `OPENAI_API_KEY` | Yes | Embeddings + image descriptions + voice transcription | | `CLAUDE_CODE_OAUTH_TOKEN` | One of these two | Claude subscription token from `claude setup-token` | | `ANTHROPIC_API_KEY` | One of these two | Anthropic API key (`sk-ant-…`), pay-as-you-go | -| `CLAUDE_MODEL` | Optional | Claude model to run (defaults to `claude-sonnet-4-6`) | -| `CLAUDE_EFFORT` | Optional | Effort level: `low`, `medium`, `high`, `xhigh`, or `max` (unset lets the model decide) | +| `CLAUDE_MODEL` | Optional | Claude model to run (defaults to `claude-sonnet-5-5`; also set by `/model`) | +| `CLAUDE_EFFORT` | Optional | Effort level: `low`, `medium`, `high`, `xhigh`, or `max` (defaults to `medium` on the default model, otherwise the model decides; also set by `/effort`) | | `AGENT_BACKEND` | Optional | Agent SDK to use (defaults to `claude`; set to `cursor` for Cursor SDK) | | `MEMCLAW_PLATFORM` | Optional | Front-end the bare `memclaw` launches: `telegram`, `whatsapp`, `slack`, or `terminal` (defaults to `terminal`) | | `CURSOR_API_KEY` | For Cursor backend | Cursor API key from Dashboard → Integrations | @@ -369,6 +378,7 @@ Backend choice and credentials are covered above in [Agent Backend](#agent-backe ├── MEMORY.md # Permanent / curated memories ├── AGENTS.md # Agent instructions + user preferences ├── memclaw.db # SQLite index (embeddings + FTS5) +├── sessions.json # Conversation session per chat └── memory/ ├── 2025-06-15.md # Daily notes ├── 2025-06-16.md @@ -384,7 +394,7 @@ Inspired by [OpenClaw](https://github.com/openclaw/openclaw)'s approach to AI me - **NumPy for vectors** — cosine similarity computed in-memory, no native extensions required - **Pluggable agent SDK** — Claude or Cursor drives the agentic loop that decides how to handle your input ([Agent Backend](#agent-backend)) - **Chunking with overlap** — ~300-word chunks with 60-word overlap preserve context across boundaries -- **Auto-consolidation** — daily files are periodically distilled into MEMORY.md +- **Auto-consolidation** — daily files are periodically distilled into MEMORY.md, in the background after a reply - **Filesystem guardrail** — every `file_read` / `file_write` tool call is path-checked, blocking anything outside `~/.memclaw/` - **Embedding cache** — SHA-256 content hashing skips redundant API calls diff --git a/docs/docker.md b/docs/docker.md index 53b3612..2bd5b8c 100644 --- a/docs/docker.md +++ b/docs/docker.md @@ -121,6 +121,15 @@ docker compose up -d --build as-is. No reverse proxy needed. - Memory footprint is modest (a few hundred MB), but the SQLite index grows with your memory vault, so back the volume up periodically. -- The Claude backend (`AGENT_BACKEND=claude`, the default) spawns the - Claude Code CLI per turn. The Cursor backend (`AGENT_BACKEND=cursor`) - is pure-Python and slightly leaner if you don't need Claude. +- The Claude backend (`AGENT_BACKEND=claude`, the default) keeps one + Claude Code CLI process running per active chat (closed on `/new` + and shutdown). The Cursor backend (`AGENT_BACKEND=cursor`) is + pure-Python and slightly leaner if you don't need Claude. +- Conversations resume after a restart: `sessions.json` lives in the + data volume and Claude Code's own session files in the container's + `~/.claude`. That directory isn't on the volume, so after a rebuild + (`docker compose up -d --build` recreates the container) each chat + starts a fresh conversation. Memories are unaffected. +- Telegram's `/model` and `/effort` save their choice to + `~/.memclaw/.env`. A `CLAUDE_MODEL` / `CLAUDE_EFFORT` passed in from + the compose `.env` wins over that file on the next start. diff --git a/docs/memory-architecture.md b/docs/memory-architecture.md index a5a148a..575d8c9 100644 --- a/docs/memory-architecture.md +++ b/docs/memory-architecture.md @@ -6,18 +6,21 @@ Memclaw is a personal memory assistant that stores everything you tell it and re ## Short-Term Memory -Short-term memory covers anything the agent needs to remember within a single session but that doesn't persist across restarts. +Short-term memory is the conversation itself: what you and the agent said in the current chat. -### Conversation History +### Conversation Sessions -The agent keeps a rolling buffer of recent messages — both yours and its own responses. This lets it follow multi-turn conversations without losing the thread. +Each chat's conversation lives in the agent backend's own session (for Claude, a Claude Code session), not in a buffer Memclaw re-sends with every message. **How it works:** -- Each time you send a message, it's appended to an in-memory list along with a timestamp. -- When the agent responds, that response is appended too. -- After each append, the buffer is trimmed to the **union** of two rolling windows: the last 10 exchanges (20 entries) and every entry timestamped within the last hour. Whichever set is larger wins; older entries that fall outside both are dropped. -- The full history is formatted and included in the system prompt so the agent can see what was said earlier in the session. +- Every chat gets a session key: `:` for the bots (e.g. `telegram:12345`, `slack:C01ABC`), `cli` for the interactive terminal. +- The Claude backend keeps one connected Claude CLI process per chat and reuses it for every turn, so a message doesn't pay the CLI start-up cost. Turns in the same chat are serialized. +- The session id is saved in `~/.memclaw/sessions.json` (written atomically). After a restart, the next message in that chat resumes the same session, so the conversation carries on. +- If a session can't be resumed (deleted, created on another machine, ...), a warning is logged and a fresh session starts. +- Each saved session records a **fingerprint** of the code-level instructions (the system-prompt template and the backend). Claude Code keeps a session's original system prompt when it resumes, so a session created under different instructions (e.g. after an upgrade that changed the template) is not resumed; the chat starts fresh. +- `/new` (Telegram command; a message of exactly `/new` in Slack, WhatsApp and the terminal) closes that chat's session and forgets its id. Your memories are untouched. +- `/model` and `/effort` (Telegram) switch the model or effort level. Live sessions reconnect on their next message, resuming the same session with the new settings, so the conversation is kept. **What this means in practice:** @@ -28,28 +31,32 @@ You: What's my dog called? Assistant: Your dog's name is Max. ``` -Without this buffer, the second question would get "I don't know" because each message was previously independent. +The conversation is bounded only by the backend's own context management (Claude Code compacts long sessions on its own). Use `/new` to start over. -The union policy matters in two regimes. During a quick burst of chatter (say, 30 messages in five minutes) the time window keeps **all** of them in scope, even though the 10-pair floor alone would have dropped the oldest. During a slow conversation with long gaps, the time window may be empty, but the 10-pair floor still guarantees continuity across resumption. +The Cursor backend has no native sessions yet: as an interim measure it keeps the last 10 messages per chat in memory and folds them into each prompt. That transcript does not survive a restart. -**Limitations:** +### System Prompt and Per-Message Context -- The buffer resets when the process restarts. It is not written to disk. -- For images, the history stores `[User sent a photo]` as a placeholder — not the actual image data. -- Both bounds are configurable: `conversation_history_limit` (default `10` pairs) and `conversation_history_window_minutes` (default `60`). +The system prompt is **static for a session**, so the backend can prompt-cache it. It is built when a session is created and contains: -### Memory Context Injection +1. **AGENTS.md** — your behavioural instructions (editable via the `update_instructions` tool). +2. Reply-formatting rules and a description of the per-message context block. +3. **MEMORY.md** — your permanent memory (see Permanent Memory in the Prompt below). -Before every response, the agent builds a context block from long-term storage and injects it into the system prompt. This is ephemeral — it's assembled fresh for each message based on what seems relevant. +Everything that changes per turn goes at the **top of the user message**, in a compact `` block: -The context has two parts: +- The current local date and time (used for dates and reminders). +- **Relevant memories** — the top 5 search results for the message (hybrid search), each cut to 500 characters. MEMORY.md hits are skipped since MEMORY.md is already in the system prompt. +- **Delivered reminders** — reminders that fired in this chat since the last message, so the agent has context if you reply to one. +- **Updated AGENTS.md / MEMORY.md** — see below. -1. **Permanent memory (MEMORY.md)** — always included (see Smart Context Strategy below). -2. **Relevant memories** — the top 10 search results from the full memory index, matched against the current message using hybrid search. +The block stays in the session transcript, which is why it is kept small. -This means the agent always has access to your key facts and any memories related to what you're currently talking about, even if you didn't explicitly ask to search. +### Keeping AGENTS.md and MEMORY.md Current ---- +The session stores a hash of the AGENTS.md and MEMORY.md content it has seen (persisted in `sessions.json`). Before each message the files are hashed again; if either changed (an `update_instructions` call, a consolidation, a manual edit), its new content is included **once** in that message's context block, labelled as superseding the earlier version, and the stored hash is updated. The conversation continues without a reset. Changes the session makes itself during a turn (`memory_save` with `permanent=true`, `update_instructions`) count as seen and aren't sent back; a consolidation that rewrites MEMORY.md during a turn is. + +Because Claude Code keeps a resumed session's original system prompt, the stored hashes always describe what the session has actually seen, including across restarts. ## Long-Term Memory @@ -62,6 +69,7 @@ Long-term memory is everything that survives between sessions. It lives on disk ├── MEMORY.md # Curated permanent knowledge base ├── memclaw.db # SQLite: chunks, embeddings, FTS index, cache ├── meta.json # Consolidation tracking metadata +├── sessions.json # Per-chat conversation session ids + hashes └── memory/ ├── 2025-03-01.md # Daily log for March 1 ├── 2025-03-02.md # Daily log for March 2 @@ -88,27 +96,22 @@ Daily files accumulate over time. Consolidation is a periodic process that disti **When it runs:** -- Automatically, before each message, if 7 or more daily files haven't been consolidated yet. +- Automatically, in the background after each reply, if 7 or more daily files haven't been consolidated yet. At most one consolidation runs at a time; failures are logged and retried after a later message. - Manually, via the `memclaw consolidate` CLI command (which forces it regardless of count). **What it does:** 1. Collects all unconsolidated daily files (up to 30,000 characters). 2. Reads the current MEMORY.md. -3. Sends both to Claude with instructions to extract durable facts, merge with existing content, remove outdated entries, and output a clean structured document. +3. Sends both to the agent backend (a short-lived, tool-free call) with instructions to extract durable facts, merge with existing content, remove outdated entries, and output a clean structured document. 4. Overwrites MEMORY.md with the result. 5. Records the latest consolidated date in `meta.json` so those files aren't processed again. The result is a MEMORY.md organized into sections like Preferences, Projects, People, Key Facts, and Decisions — with the most important information at the top. -### Smart Context Strategy - -When MEMORY.md is included in the system prompt, it uses an adaptive strategy rather than a hard character cutoff: - -- **Small file (under 4,000 characters):** Included in full. Nothing is lost. -- **Large file (over 4,000 characters):** The first 2,000 characters are always included (this is where consolidation places the most important content). Then a semantic search runs against just the MEMORY.md chunks using the current message as query, and the top 3 matching sections are appended. +### Permanent Memory in the Prompt -This means a 10,000-character MEMORY.md doesn't lose relevant information — it just includes the top section plus the parts most related to what you're currently asking about. +MEMORY.md goes into the session's system prompt in full (consolidation keeps it under ~5,000 characters; anything past 20,000 characters is cut with a note pointing the agent to `memory_search`). Since the system prompt is cached, carrying the whole file costs little. When consolidation rewrites MEMORY.md mid-conversation, the new version reaches the session through the context block as described above. --- @@ -157,7 +160,7 @@ This means searching "pizza" returns diverse results (your preference, a restaur ### File Filtering -Search can be restricted to a specific file by passing a `file_filter` string. Internally this is used to search only within MEMORY.md when building the smart context strategy. It filters results by checking if the file path contains the given substring. +Search can be restricted to a specific file by passing a `file_filter` string. It filters results by checking if the file path contains the given substring. ### Full Search Pipeline @@ -235,8 +238,6 @@ All memory-related settings live in `MemclawConfig`: | `text_weight` | `0.3` | Weight for keyword search in hybrid merge | | `decay_half_life_days` | `30` | Days until a daily memory's score halves (0 = disabled) | | `mmr_lambda` | `0.7` | Relevance vs. diversity trade-off in MMR (1.0 = pure relevance) | -| `conversation_history_limit` | `10` | Minimum message pairs to keep in session buffer (count floor of the union policy) | -| `conversation_history_window_minutes` | `60` | Time window for the session buffer — every entry newer than this is also retained | | `consolidation_threshold` | `7` | Number of unconsolidated daily files before auto-consolidation | --- @@ -244,22 +245,23 @@ All memory-related settings live in `MemclawConfig`: ## Data Flow Summary ``` -User Message +User Message (chat → session key) │ -├─ 1. Consolidation check (merge daily files → MEMORY.md if threshold reached) -├─ 2. Append to conversation history buffer -├─ 3. Build context: -│ ├─ Include MEMORY.md (full or smart-truncated) -│ └─ Semantic search for 10 relevant memory chunks -├─ 4. Format history for system prompt -├─ 5. Send to Claude with context + history + tools +├─ 1. Look up the chat's saved session (sessions.json; dropped if the fingerprint changed) +├─ 2. Build the block: +│ ├─ Current local time +│ ├─ Reminders delivered since the last message +│ ├─ AGENTS.md / MEMORY.md, if changed since the session last saw them +│ └─ Semantic search for 5 relevant memory chunks +├─ 3. Run the turn in the chat's session (reuse the live client, or resume / +│ start one with the static system prompt: AGENTS.md + rules + MEMORY.md) │ └─ Agent may call: │ ├─ memory_save → write to daily file or MEMORY.md → index │ ├─ memory_search → run full search pipeline -│ ├─ image_save / telegram_image_save → store image metadata +│ ├─ image_save → store image metadata │ └─ image_search → vector search over image descriptions -├─ 6. Append assistant response to history buffer -├─ 7. Trim history to union of (last N pairs) and (entries within the time window) +├─ 4. Save the session id + AGENTS.md / MEMORY.md hashes +├─ 5. Schedule consolidation in the background (if threshold reached) │ -└─ Return response (+ any found images for Telegram) +└─ Return response (+ any found images) ``` diff --git a/memclaw/agent.py b/memclaw/agent.py index c2e8cfc..2c57df2 100644 --- a/memclaw/agent.py +++ b/memclaw/agent.py @@ -1,17 +1,24 @@ """Memclaw agent — backend-agnostic orchestration over a pluggable agent SDK. -The agent owns memory, search, history, consolidation, and the system-prompt -shape, but delegates every LLM call to an `AgentBackend` (see -`memclaw.backends`). The backend is selected by `config.agent_backend` -(env var `AGENT_BACKEND`), defaulting to `claude`. +The agent owns memory, search, consolidation, the system-prompt shape and +the per-chat session bookkeeping, but delegates every LLM call to an +`AgentBackend` (see `memclaw.backends`). The backend is selected by +`config.agent_backend` (env var `AGENT_BACKEND`), defaulting to `claude`. + +Each chat's conversation lives in the backend's own session. The system +prompt is static per session (so it can be prompt-cached); everything that +changes per turn goes into a `` block at the top of the user +message instead. """ from __future__ import annotations import asyncio +import contextlib +import hashlib import json import time -from datetime import date, datetime, timedelta +from datetime import date, datetime from pathlib import Path from loguru import logger @@ -21,15 +28,15 @@ from .index import MemoryIndex from .reminders import ReminderScheduler from .search import HybridSearch +from .sessions import SessionStore from .store import MemoryStore from .tools import ToolExecutor # ── Prompts ────────────────────────────────────────────────────────── +# Static for the whole session. Changing this template changes the session +# fingerprint, so existing sessions start fresh instead of being resumed. _SYSTEM_PROMPT_TEMPLATE = """\ -Today's date: {today} -Current local time: {now} - {agent_instructions} === REPLY FORMATTING === @@ -44,11 +51,17 @@ - No backticks or fenced code blocks. - No `[label](url)` links — write the bare URL. -=== MEMORY CONTEXT === -{context} +=== PER-MESSAGE CONTEXT === +Each user message starts with a block added by Memclaw, not typed \ +by the user. It holds the current local time (use it for dates and \ +reminders), memories that may be relevant to that message, reminders that \ +were delivered to the user since their last message, and the new content of \ +AGENTS.md or MEMORY.md when either file changes — such an update supersedes \ +the version in this system prompt and any earlier update. Don't mention the \ +block itself to the user. -=== CONVERSATION HISTORY === -{history} +=== PERMANENT MEMORY (MEMORY.md) === +{permanent_memory} IMPORTANT: When the user gives you a behavioural instruction (e.g. "always respond \ in Spanish", "be more formal", "never use emojis"), you MUST call the \ @@ -56,6 +69,18 @@ future conversation. """ +# MEMORY.md is meant to stay under ~5,000 characters (see consolidation); +# this cap only guards against a runaway file. +_MEMORY_PROMPT_CHARS = 20000 + +# Per-turn memory search results. Kept small because every context block +# stays in the session transcript. +_CONTEXT_RESULTS = 5 +_CONTEXT_SNIPPET_CHARS = 500 + +# Session key for the interactive terminal (no chat id). +_CLI_SESSION_KEY = "cli" + def _load_agent_instructions(config: MemclawConfig) -> str: agent_file = config.agent_file @@ -64,6 +89,30 @@ def _load_agent_instructions(config: MemclawConfig) -> str: return "You are Memclaw, a personal memory assistant." +def _load_permanent_memory(config: MemclawConfig) -> str: + memory = config.memory_file.read_text().strip() if config.memory_file.exists() else "" + if not memory: + return "(empty)" + if len(memory) > _MEMORY_PROMPT_CHARS: + return memory[:_MEMORY_PROMPT_CHARS] + "\n\n[... truncated; use memory_search for more]" + return memory + + +def _hash(text: str) -> str: + return hashlib.sha256(text.encode()).hexdigest()[:16] + + +def instructions_fingerprint(backend_name: str) -> str: + """Hash of the code-level instructions a session is started with. + + Claude Code keeps a session's original system prompt when it is resumed, + so a session started under a different template (or backend) must not be + resumed. AGENTS.md / MEMORY.md are tracked separately per session and + sent as updates instead, so editing them doesn't reset conversations. + """ + return _hash(f"{backend_name}\0{_SYSTEM_PROMPT_TEMPLATE}") + + _CONSOLIDATION_PROMPT = """\ You are a memory consolidation assistant. Your job is to distill daily memory \ logs into a curated, permanent knowledge base. @@ -101,7 +150,7 @@ def _load_agent_instructions(config: MemclawConfig) -> str: class MemclawAgent: """Unified agent for both interactive CLI and messaging bots. - Memory, search, history, and consolidation orchestration live here. + Memory, search, consolidation, and session orchestration live here. Every LLM call is delegated to a pluggable `AgentBackend` chosen via `config.agent_backend`. """ @@ -121,7 +170,6 @@ def __init__( self.search = HybridSearch(config, self.index) self.scheduler = scheduler self._found_images: list[dict] = [] - self._history: list[dict] = [] self._tools = ToolExecutor( config=config, store=self.store, @@ -133,43 +181,41 @@ def __init__( ) self.backend: AgentBackend = backend or build_backend(config) self._backend_started = False + self.sessions = SessionStore(config.memory_dir / "sessions.json") + self._fingerprint = instructions_fingerprint(self.backend.name) + # Delivered reminders not yet seen by the chat's session. + self._pending_notes: dict[str, list[str]] = {} + # Bumped by /new so a turn that was running doesn't re-save the + # session it just reset. + self._generations: dict[str, int] = {} + self._consolidation_task: asyncio.Task | None = None + # Bumped whenever consolidation rewrites MEMORY.md, so a turn can tell + # its own MEMORY.md writes from a rewrite it hasn't seen. + self._memory_rewrites = 0 + + # ── Sessions ───────────────────────────────────────────────────── + + def _session_key(self, chat_id: str | None) -> str: + if chat_id is None: + return _CLI_SESSION_KEY + return f"{self.platform}:{chat_id}" if self.platform else str(chat_id) + + def record_reminder_fired(self, text: str, chat_id: str | None = None): + """Remember a delivered reminder so the chat's next message carries + it, giving the agent context if the user replies to it.""" + self._pending_notes.setdefault(self._session_key(chat_id), []).append(text) + + async def reset_conversation(self, chat_id: str | None = None) -> None: + """Start a fresh conversation for *chat_id* (the `/new` command).""" + key = self._session_key(chat_id) + self._generations[key] = self._generations.get(key, 0) + 1 + await self.backend.reset_session(key) + self.sessions.drop(key) + self._pending_notes.pop(key, None) + logger.info("Started a new conversation for {k}", k=key) # ── Startup / sync ─────────────────────────────────────────────── - def record_reminder_fired(self, text: str): - """Append a delivered reminder to history so the agent has context - if the user replies to it.""" - self._history.append({ - "role": "assistant", - "content": f"[Reminder fired] {text}", - "timestamp": datetime.now().isoformat(), - }) - self._trim_history() - - def _trim_history(self) -> None: - """Retain the union of: the last N pairs, and entries newer than the - configured time window. Entries are appended chronologically, so the - "within window" subset is always a contiguous suffix.""" - if not self._history: - return - count_floor = self.config.conversation_history_limit * 2 - cutoff = datetime.now() - timedelta( - minutes=self.config.conversation_history_window_minutes - ) - in_window = 0 - for entry in reversed(self._history): - try: - ts = datetime.fromisoformat(entry["timestamp"]) - except (KeyError, ValueError): - break - if ts >= cutoff: - in_window += 1 - else: - break - keep = max(count_floor, in_window) - if len(self._history) > keep: - self._history = self._history[-keep:] - async def start(self, *, include_backend: bool = True): await self.index.sync() if include_backend: @@ -177,7 +223,13 @@ async def start(self, *, include_backend: bool = True): self._backend_started = True async def aclose(self): - """Async shutdown: release backend resources then sync cleanup.""" + """Async shutdown: stop consolidation, release backend resources, + then sync cleanup.""" + task = self._consolidation_task + if task is not None and not task.done(): + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task if self._backend_started: await self.backend.on_agent_shutdown() self._backend_started = False @@ -258,6 +310,7 @@ async def _maybe_consolidate( return False self.config.memory_file.write_text(result_text) + self._memory_rewrites += 1 last_date_str = unconsolidated[-1].stem try: @@ -282,34 +335,64 @@ async def _maybe_consolidate( ) return True + def _schedule_consolidation(self) -> None: + """Run the consolidation check in the background, after the reply, + so it never delays one. At most one runs at a time.""" + if self._consolidation_task is not None and not self._consolidation_task.done(): + return + self._consolidation_task = asyncio.create_task(self._consolidate_in_background()) + + async def _consolidate_in_background(self) -> None: + try: + await self._maybe_consolidate() + except Exception as exc: + logger.exception("Background consolidation failed: {exc}", exc=exc) + # ── Context builder ────────────────────────────────────────────── - async def build_context(self, message: str) -> str: - parts: list[str] = [] - - memory_content = self.store.read_file(self.config.memory_file) - if memory_content.strip(): - parts.append("### Permanent Memory") - if len(memory_content) <= 4000: - parts.append(memory_content) - else: - parts.append(memory_content[:2000]) - memory_results = await self.search.search( - message, limit=3, file_filter="MEMORY.md" - ) - if memory_results: - parts.append("\n#### Relevant Permanent Memory Sections") - for r in memory_results: - parts.append(r.content.strip()) - - results = await self.search.search(message, limit=10) + async def build_context( + self, + message: str, + *, + notes: list[str] | tuple[str, ...] = (), + agents_update: str | None = None, + memory_update: str | None = None, + ) -> str: + """The per-turn `` block prefixed to the user message. + + It stays in the session transcript, so it only carries what changes + per turn and keeps memory snippets short. MEMORY.md is already in + the system prompt, so its search hits are skipped. + """ + now = datetime.now() + parts = [ + "", + f"Current local time: {now:%Y-%m-%d %H:%M} ({now:%A})", + ] + if notes: + parts.append("Reminders delivered since the user's last message:") + parts.extend(f"- {note}" for note in notes) + if agents_update is not None: + parts.append("Updated AGENTS.md (supersedes the earlier version):") + parts.append(agents_update) + if memory_update is not None: + parts.append("Updated permanent memory, MEMORY.md (supersedes the earlier version):") + parts.append(memory_update) + + results = await self.search.search(message, limit=_CONTEXT_RESULTS + 2) + results = [ + r for r in results if Path(r.file_path).name != self.config.memory_file.name + ][:_CONTEXT_RESULTS] if results: - parts.append("\n### Relevant Memories") + parts.append("Relevant memories:") for r in results: - source = Path(r.file_path).stem - parts.append(f"[{source}] {r.content.strip()}") + snippet = " ".join(r.content.split()) + if len(snippet) > _CONTEXT_SNIPPET_CHARS: + snippet = snippet[:_CONTEXT_SNIPPET_CHARS] + "…" + parts.append(f"- [{Path(r.file_path).stem}] {snippet}") - return "\n\n".join(parts) if parts else "No memories found yet." + parts.append("") + return "\n".join(parts) # ── Main entry point ───────────────────────────────────────────── @@ -323,45 +406,46 @@ async def handle( ) -> tuple[str, list[dict]]: self._found_images.clear() self._tools.chat_id = chat_id - - try: - await self._maybe_consolidate() - except Exception as exc: - logger.warning("Consolidation check failed: {exc}", exc=exc) - - history_content = "[User sent a photo]" if image_b64 else message - self._history.append({ - "role": "user", - "content": history_content, - "timestamp": datetime.now().isoformat(), - }) - - context = await self.build_context(message) - - history_snapshot = self._history[:-1] - if history_snapshot: - history_lines = [] - for entry in history_snapshot: - role = "User" if entry["role"] == "user" else "Assistant" - history_lines.append(f"{role}: {entry['content']}") - history_text = "\n".join(history_lines) - else: - history_text = "(no prior messages)" + key = self._session_key(chat_id) + generation = self._generations.get(key, 0) + memory_rewrites = self._memory_rewrites agent_instructions = _load_agent_instructions(self.config) + permanent_memory = _load_permanent_memory(self.config) + agents_hash = _hash(agent_instructions) + memory_hash = _hash(permanent_memory) + + # A resumed (or live) session has the AGENTS.md / MEMORY.md it last + # saw; send whichever changed since, once. + entry = self.sessions.get(key, self._fingerprint) + notes = list(self._pending_notes.get(key, ())) + context = await self.build_context( + message, + notes=notes, + agents_update=( + agent_instructions + if entry and entry.get("agents_hash") != agents_hash else None + ), + memory_update=( + permanent_memory + if entry and entry.get("memory_hash") != memory_hash else None + ), + ) + + # Only used when the backend starts a new session. system_prompt = _SYSTEM_PROMPT_TEMPLATE.format( - today=date.today().isoformat(), - now=datetime.now().replace(microsecond=0).isoformat(), agent_instructions=agent_instructions, - context=context, - history=history_text, + permanent_memory=permanent_memory, ) t0 = time.perf_counter() result = await self.backend.run_turn( system_prompt=system_prompt, + context=context, user_message=message, tool_executor=self._tools, + session_key=key, + resume_session_id=entry["session_id"] if entry else None, image_b64=image_b64, image_media_type=image_media_type, max_turns=10, @@ -386,21 +470,33 @@ async def handle( turns=result.num_turns or 1, ms=elapsed_ms, tokens=token_summary, ) - response_text = result.text or "I couldn't generate a response." - self._history.append({ - "role": "assistant", - "content": response_text, - "timestamp": datetime.now().isoformat(), - }) - - self._trim_history() + # The session has now seen these notes and these file versions. Its + # own writes during the turn (memory_save permanent=true, + # update_instructions) count as seen too, so they aren't sent back + # next turn; a consolidation rewrite during the turn doesn't. + agents_hash = _hash(_load_agent_instructions(self.config)) + if self._memory_rewrites == memory_rewrites: + memory_hash = _hash(_load_permanent_memory(self.config)) + pending = self._pending_notes.get(key) + if pending: + del pending[:len(notes)] + if result.session_id and self._generations.get(key, 0) == generation: + self.sessions.set(key, { + "session_id": result.session_id, + "fingerprint": self._fingerprint, + "agents_hash": agents_hash, + "memory_hash": memory_hash, + }) + + self._schedule_consolidation() + response_text = result.text or "I couldn't generate a response." return (response_text, list(self._found_images)) def _close_sync_only(self) -> None: - task = getattr(self, "_sync_task", None) - if task is not None and not task.done(): - task.cancel() + for task in (getattr(self, "_sync_task", None), self._consolidation_task): + if task is not None and not task.done(): + task.cancel() self.index.close() def close(self): diff --git a/memclaw/backends/base.py b/memclaw/backends/base.py index aed9466..5a28f93 100644 --- a/memclaw/backends/base.py +++ b/memclaw/backends/base.py @@ -2,9 +2,14 @@ A backend wraps a specific agent SDK (Claude Agent SDK, Cursor SDK, OpenAI Agents SDK, ...) and exposes a uniform surface to MemclawAgent. The protocol -is intentionally narrow: it covers exactly the two interactions MemclawAgent -needs — one-shot text generation (used for memory consolidation) and a full -agentic turn with tool access (used for every user message). +is intentionally narrow: it covers one-shot text generation (used for memory +consolidation), a full agentic turn with tool access (used for every user +message), per-chat conversation sessions, and model / effort selection. + +Conversations live in the backend's own session, keyed by a ``session_key`` +MemclawAgent derives from the chat. MemclawAgent persists the session id a +turn reports and hands it back as ``resume_session_id``, so a restart picks +the conversation up where it left off. Adding a new backend means implementing this protocol and registering the class in `memclaw.backends.__init__.REGISTRY`. Nothing else in the project @@ -22,6 +27,7 @@ class in `memclaw.backends.__init__.REGISTRY`. Nothing else in the project from ..config import MemclawConfig from ..tools import ToolExecutor + from .claude_models import ModelInfo @dataclass @@ -41,6 +47,9 @@ class TurnResult: cache_read_tokens: int = 0 cache_creation_tokens: int = 0 cost_usd: float | None = None + # The backend's native session id for this conversation, or None when the + # backend has no resumable sessions. + session_id: str | None = None @runtime_checkable @@ -56,6 +65,9 @@ class AgentBackend(Protocol): # a subscription/plan. bills_per_token: bool + # Whether /model and /effort can switch this backend's model at runtime. + supports_model_selection: ClassVar[bool] + def __init__(self, config: "MemclawConfig") -> None: ... # Configuration -------------------------------------------------------- @@ -135,13 +147,24 @@ async def run_turn( self, *, system_prompt: str, + context: str, user_message: str, tool_executor: "ToolExecutor", + session_key: str, + resume_session_id: str | None = None, image_b64: str | None = None, image_media_type: str = "image/jpeg", max_turns: int = 10, ) -> TurnResult: - """Run one full agentic turn with tool access. + """Run one full agentic turn with tool access, in the conversation + identified by *session_key*. + + ``system_prompt`` is static per session: it is only used when the + backend starts a new session. ``context`` is the per-turn context + block (time, relevant memories, ...) that must reach the model with + this message, ahead of ``user_message``. ``resume_session_id`` is + the session to resume if the backend has to (re)connect; if it can't + be resumed the backend starts a fresh one. The backend is responsible for translating ``TOOL_DEFINITIONS`` into its SDK's tool format and routing tool calls back through @@ -153,6 +176,39 @@ async def run_turn( """ ... + async def reset_session(self, session_key: str) -> None: + """Forget the conversation for *session_key* (the `/new` command).""" + ... + + # Model selection (only used when supports_model_selection) ------------ + @property + def model(self) -> str: + """The model the next turn will run on.""" + ... + + @property + def effort(self) -> str | None: + """The effort level the next turn will run with (None: model default).""" + ... + + async def list_models(self) -> list["ModelInfo"]: + """Models the configured credential can use. Raises on failure.""" + ... + + async def effort_levels(self) -> list[str] | None: + """Effort levels the current model accepts ([]: none, None: unknown).""" + ... + + async def set_model(self, model_id: str) -> str: + """Switch model, keeping conversations. Returns a note for the user + (e.g. when the effort level had to change). Raises ValueError for + an unknown model.""" + ... + + async def set_effort(self, level: str) -> None: + """Set the effort level. Raises ValueError if the model rejects it.""" + ... + # Subclasses register themselves through this attribute name; see # memclaw/backends/__init__.py. diff --git a/memclaw/backends/claude.py b/memclaw/backends/claude.py index 188c15b..03cdb2c 100644 --- a/memclaw/backends/claude.py +++ b/memclaw/backends/claude.py @@ -2,8 +2,12 @@ Wraps `claude-agent-sdk` (which itself shells out to the Claude CLI) into the neutral `AgentBackend` protocol. All Claude-specific glue — env scrubbing, -MCP server wrapping for tools, stream-json image protocol, response parsing -— lives in this file. +MCP server wrapping for tools, stream-json image protocol, response parsing, +per-chat CLI sessions — lives in this file. + +Each chat keeps one connected `ClaudeSDKClient` (one Claude CLI process) +that is reused across turns, so the conversation lives in the CLI's own +session instead of being re-sent with every message. """ from __future__ import annotations @@ -11,6 +15,7 @@ import asyncio import json import os +from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any, AsyncIterator, ClassVar @@ -30,7 +35,6 @@ from .base import TurnResult from .claude_models import SDK_EFFORT_LEVELS, ModelInfo, fetch_models from .mcp_tools import MCP_SERVER_NAME -from .tool_policy import BUILTIN_TOOLS_DISALLOW if TYPE_CHECKING: from rich.console import Console @@ -39,11 +43,11 @@ from ..tools import ToolExecutor -_DEFAULT_MODEL = "claude-sonnet-4-6" +_DEFAULT_MODEL = "claude-sonnet-5-5" -# Preselected in the wizard for any model that supports effort at all. -# Every model that reports effort support accepts this level. -_DEFAULT_EFFORT = "high" +# Used when CLAUDE_EFFORT is unset and the default model runs, and +# preselected in the wizard for any model that supports effort at all. +_DEFAULT_EFFORT = "medium" # Anthropic list prices, (input, output) USD per 1M tokens. Only used as a # fallback when the SDK doesn't return total_cost_usd for an API-key turn. @@ -105,16 +109,40 @@ def _prices_per_m(model: str) -> tuple[float, float]: def _resolve_effort(config: "MemclawConfig") -> str | None: """The configured effort level, or None to let the model decide. - None makes the SDK omit the flag entirely, so an unset level costs - nothing. Anything the SDK does not recognise is treated as unset: - ~/.memclaw/.env is a plain text file, and a typo there should not reach - the CLI as an argument it will reject. + Anything the SDK does not recognise is treated as unset: ~/.memclaw/.env + is a plain text file, and a typo there should not reach the CLI as an + argument it will reject. An unset level means `_DEFAULT_EFFORT` on the + default model and None (the SDK omits the flag) on any other model. Whether the chosen *model* accepts an effort level is a separate question, - settled in the wizard — it stores no level for a model that reports none. + settled in the wizard and by /model — they store no level for a model + that reports none. """ effort = (config.claude_effort or "").strip().lower() - return effort if effort in SDK_EFFORT_LEVELS else None + if effort in SDK_EFFORT_LEVELS: + return effort + return _DEFAULT_EFFORT if _resolve_model(config) == _DEFAULT_MODEL else None + + +def _fit_effort(levels: list[str], preferred: str) -> str | None: + """The level to use on a model supporting *levels* when *preferred* is + wanted: *preferred* itself, else the nearest supported level below it + (else the lowest). None when the model takes no effort at all.""" + if not levels: + return None + if preferred in levels: + return preferred + order = list(SDK_EFFORT_LEVELS) + rank = order.index(preferred) if preferred in order else len(order) + below = [lv for lv in levels if lv in order and order.index(lv) < rank] + return below[-1] if below else levels[0] + + +def _credential(config: "MemclawConfig") -> str: + """The credential matching `_claude_auth_mode`.""" + if _claude_auth_mode(config) == "subscription": + return config.claude_code_oauth_token + return config.anthropic_api_key def _build_env(config: "MemclawConfig") -> dict[str, str]: @@ -174,6 +202,12 @@ async def wrapper(args: dict[str, Any]) -> dict[str, Any]: # ── Image-input streaming protocol ────────────────────────────────── +def _compose_prompt(context: str, user_message: str) -> str: + """The per-turn context block goes at the top of the user message, so + the session's system prompt stays static and cacheable.""" + return f"{context}\n\n{user_message}" if context else user_message + + async def _image_prompt_stream( message: str, image_b64: str, image_media_type: str, ) -> AsyncIterator[dict[str, Any]]: @@ -205,11 +239,24 @@ async def _image_prompt_stream( # ── Backend ───────────────────────────────────────────────────────── +@dataclass +class _Session: + """One chat's live Claude CLI connection.""" + + # Serializes turns (and /new) within one chat. + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + client: ClaudeSDKClient | None = None + # (model, effort) the client was started with. After /model or /effort + # the next turn reconnects, resuming the same session with the new ones. + options_key: tuple[str, str | None] | None = None + + class ClaudeAgentBackend: """Claude Agent SDK implementation of the AgentBackend protocol.""" name: ClassVar[str] = "claude" display_name: ClassVar[str] = "Claude Agent SDK" + supports_model_selection: ClassVar[bool] = True def __init__(self, config: "MemclawConfig") -> None: self.config = config @@ -217,6 +264,11 @@ def __init__(self, config: "MemclawConfig") -> None: self._model = _resolve_model(config) self._effort = _resolve_effort(config) self._mcp_server: Any = None # built lazily, needs a ToolExecutor + self._sessions: dict[str, _Session] = {} + self._models: list[ModelInfo] = [] # filled by list_models() + # Last total_cost_usd seen per session id. The CLI reports a running + # total for the session (carried over on resume), not a per-turn cost. + self._session_costs: dict[str, float] = {} # Subscription billing → no per-message dollar figure in logs. self.bills_per_token = _claude_auth_mode(config) == "api_key" @@ -399,7 +451,7 @@ def _pick_effort( """Show the effort picker for a model that supports one. *levels* only ever holds levels this model accepts, so the preselected - value falls back to the first one when the model has no `high`. + value falls back to the first one when the model has no `medium`. """ preferred = current_effort if current_effort in levels else _DEFAULT_EFFORT index = choose( @@ -421,7 +473,164 @@ async def on_agent_start(self, tool_executor: "ToolExecutor") -> None: pass async def on_agent_shutdown(self) -> None: - pass + for session in self._sessions.values(): + await self._disconnect(session) + self._sessions.clear() + + # -- Sessions ---------------------------------------------------------- + + def _session_options( + self, *, system_prompt: str, resume: str | None, max_turns: int, + ) -> ClaudeAgentOptions: + return ClaudeAgentOptions( + env=self._env, + model=self._model, + effort=self._effort, + system_prompt=system_prompt, + # Claude Code stores sessions per working directory; a fixed one + # lets a restarted process find (and resume) them. + cwd=str(self.config.memory_dir), + # Don't load ~/.claude settings, plugins, CLAUDE.md or the user's + # MCP servers: faster startup and nothing outside our control. + setting_sources=[], + strict_mcp_config=True, + mcp_servers={MCP_SERVER_NAME: self._mcp_server}, + # No built-in tools at all; only our in-process mcp__memclaw__*. + tools=[], + allowed_tools=_ALLOWED_TOOLS, + # Anything not in allowed_tools is denied, never prompted for. + permission_mode="dontAsk", + # No @file expansion / slash commands from chat text. + verbatim_prompts=True, + max_turns=max_turns, + resume=resume, + stderr=lambda line: logger.debug("claude-cli: {line}", line=line.rstrip()), + ) + + async def _connect( + self, *, system_prompt: str, resume: str | None, max_turns: int, + ) -> ClaudeSDKClient: + """Start a CLI for one chat, resuming *resume* when given. A session + that can't be resumed (deleted, from another machine, ...) is + replaced by a fresh one rather than failing the turn.""" + if resume: + client = ClaudeSDKClient(options=self._session_options( + system_prompt=system_prompt, resume=resume, max_turns=max_turns, + )) + try: + await client.connect() + return client + except Exception as exc: + logger.warning( + "Could not resume Claude session {id} ({exc}); starting a fresh one", + id=resume, exc=exc, + ) + client = ClaudeSDKClient(options=self._session_options( + system_prompt=system_prompt, resume=None, max_turns=max_turns, + )) + await client.connect() + return client + + @staticmethod + async def _disconnect(session: _Session) -> None: + client, session.client = session.client, None + if client is None: + return + try: + await client.disconnect() + except Exception as exc: + logger.warning("Error while disconnecting Claude CLI: {exc}", exc=exc) + + async def reset_session(self, session_key: str) -> None: + session = self._sessions.get(session_key) + if session is None: + return + async with session.lock: + await self._disconnect(session) + + # -- Model selection --------------------------------------------------- + + @property + def model(self) -> str: + return self._model + + @property + def effort(self) -> str | None: + return self._effort + + async def list_models(self) -> list[ModelInfo]: + self._models = await fetch_models(_claude_auth_mode(self.config), _credential(self.config)) + return self._models + + async def _model_info(self, model_id: str) -> ModelInfo | None: + """*model_id*'s entry in the model list, or None if it isn't listed + or the list can't be fetched.""" + if not self._models: + try: + await self.list_models() + except Exception as exc: + logger.warning("Could not fetch the Claude model list: {exc}", exc=exc) + return None + return next((m for m in self._models if m.id == model_id), None) + + async def effort_levels(self) -> list[str] | None: + info = await self._model_info(self._model) + return list(info.effort_levels) if info else None + + async def set_model(self, model_id: str) -> str: + model_id = model_id.strip() + info = await self._model_info(model_id) + if info is None and self._models: + raise ValueError(f"Unknown model '{model_id}'.") + effort = self._effort + note = "" + if info is not None: + # Keep the effort if the new model takes it, otherwise the nearest + # level it does take (or none at all). + effort = _fit_effort(info.effort_levels, self._effort or _DEFAULT_EFFORT) + if self._effort and effort is None: + note = f"{info.display_name} doesn't support an effort setting, so none is used." + elif self._effort and effort != self._effort: + note = ( + f"{info.display_name} doesn't support '{self._effort}' effort; " + f"using '{effort}' instead." + ) + self._apply_model(model_id, effort) + return note + + async def set_effort(self, level: str) -> None: + level = level.strip().lower() + if level not in SDK_EFFORT_LEVELS: + raise ValueError( + f"Unknown effort '{level}'. Levels: {', '.join(SDK_EFFORT_LEVELS)}." + ) + levels = await self.effort_levels() + if levels == []: + raise ValueError("The current model doesn't support an effort setting.") + if levels is not None and level not in levels: + raise ValueError( + f"The current model doesn't support '{level}'. " + f"Supported: {', '.join(levels)}." + ) + self._apply_model(self._model, level) + + def _apply_model(self, model: str, effort: str | None) -> None: + """Use *model* / *effort* from the next turn on and save them to + ~/.memclaw/.env. Live chats reconnect on their next turn, resuming + their session (Claude Code resumes fine under a different model).""" + from ..setup import update_env_file # local import — setup imports backends + + logger.info( + "Claude model: {m} effort: {e} -> {m2} effort: {e2}", + m=self._model, e=self._effort, m2=model, e2=effort, + ) + self._model, self._effort = model, effort + self.config.claude_model = model + self.config.claude_effort = effort or "" + if effort: + update_env_file({"CLAUDE_MODEL": model, "CLAUDE_EFFORT": effort}) + else: + update_env_file({"CLAUDE_MODEL": model}, drop_keys=["CLAUDE_EFFORT"]) # -- Runtime --------------------------------------------------------- @@ -436,8 +645,10 @@ async def run_one_shot( model=self._model, effort=self._effort, system_prompt=system_prompt, - setting_sources=None, - disallowed_tools=BUILTIN_TOOLS_DISALLOW, + setting_sources=[], + strict_mcp_config=True, + tools=[], + verbatim_prompts=True, max_turns=1, ) @@ -458,8 +669,11 @@ async def run_turn( self, *, system_prompt: str, + context: str, user_message: str, tool_executor: "ToolExecutor", + session_key: str, + resume_session_id: str | None = None, image_b64: str | None = None, image_media_type: str = "image/jpeg", max_turns: int = 10, @@ -467,64 +681,41 @@ async def run_turn( if self._mcp_server is None: self._mcp_server = _build_mcp_server(tool_executor) - options = ClaudeAgentOptions( - env=self._env, - model=self._model, - effort=self._effort, - system_prompt=system_prompt, - setting_sources=None, - mcp_servers={MCP_SERVER_NAME: self._mcp_server}, - allowed_tools=_ALLOWED_TOOLS, - disallowed_tools=BUILTIN_TOOLS_DISALLOW, - # SAFETY: bypassPermissions is only safe because allowed_tools - # restricts execution to mcp__memclaw__* (our in-process server) - # and BUILTIN_TOOLS_DISALLOW blocks Claude Code's built-ins. - # If either guardrail is loosened, revisit this — bypass mode - # would otherwise turn any future broad tool into an RCE vector. - permission_mode="bypassPermissions", - max_turns=max_turns, - ) - - last_text = "" - result = TurnResult(text="") - - async with ClaudeSDKClient(options=options) as client: - if image_b64: - await client.query(_image_prompt_stream( - user_message, image_b64, image_media_type, - )) - else: - await client.query(user_message) - - async for msg in client.receive_response(): - if isinstance(msg, AssistantMessage): - turn_text = "" - for block in msg.content: - if isinstance(block, TextBlock): - turn_text += block.text - elif isinstance(block, ToolUseBlock): - args_str = json.dumps(block.input, ensure_ascii=False) - if len(args_str) > 300: - args_str = args_str[:300] + "..." - tool_name = block.name - prefix = f"mcp__{MCP_SERVER_NAME}__" - if tool_name.startswith(prefix): - tool_name = tool_name[len(prefix):] - logger.info("Tool call: {name}({args})", name=tool_name, args=args_str) - if turn_text: - last_text = turn_text - elif isinstance(msg, ResultMessage): - result.num_turns = msg.num_turns - result.cost_usd = msg.total_cost_usd - if msg.usage: - result.input_tokens = msg.usage.get("input_tokens", 0) or 0 - result.output_tokens = msg.usage.get("output_tokens", 0) or 0 - result.cache_read_tokens = msg.usage.get("cache_read_input_tokens", 0) or 0 - result.cache_creation_tokens = msg.usage.get("cache_creation_input_tokens", 0) or 0 - if msg.result and not last_text: - last_text = msg.result - - # Fallback cost estimate when the SDK didn't supply one. + session = self._sessions.setdefault(session_key, _Session()) + async with session.lock: + options_key = (self._model, self._effort) + if session.client is not None and session.options_key != options_key: + await self._disconnect(session) + if session.client is None: + session.client = await self._connect( + system_prompt=system_prompt, + resume=resume_session_id, + max_turns=max_turns, + ) + session.options_key = options_key + try: + result = await self._query( + session.client, + _compose_prompt(context, user_message), + image_b64=image_b64, + image_media_type=image_media_type, + ) + except BaseException: + # The CLI may be dead or mid-turn; reconnect (resuming) next time. + await self._disconnect(session) + raise + + # Turn the session's running total into this turn's cost. Without a + # baseline (first turn after a restart resumed it) estimate instead. + total = result.cost_usd + result.cost_usd = None + if total is not None and result.session_id: + previous = self._session_costs.get(result.session_id) + self._session_costs[result.session_id] = total + if previous is not None or resume_session_id != result.session_id: + result.cost_usd = max(total - (previous or 0.0), 0.0) + + # Fallback cost estimate when the SDK didn't supply a usable one. if result.cost_usd is None and self.bills_per_token: input_per_m, output_per_m = _prices_per_m(self._model) result.cost_usd = ( @@ -532,6 +723,52 @@ async def run_turn( + result.output_tokens * output_per_m + result.cache_read_tokens * input_per_m * _CACHE_READ_MULTIPLIER ) / 1_000_000 + return result + + @staticmethod + async def _query( + client: ClaudeSDKClient, + prompt: str, + *, + image_b64: str | None, + image_media_type: str, + ) -> TurnResult: + """Send one user message on a connected client and collect the reply.""" + if image_b64: + await client.query(_image_prompt_stream(prompt, image_b64, image_media_type)) + else: + await client.query(prompt) + + last_text = "" + result = TurnResult(text="") + async for msg in client.receive_response(): + if isinstance(msg, AssistantMessage): + turn_text = "" + for block in msg.content: + if isinstance(block, TextBlock): + turn_text += block.text + elif isinstance(block, ToolUseBlock): + args_str = json.dumps(block.input, ensure_ascii=False) + if len(args_str) > 300: + args_str = args_str[:300] + "..." + tool_name = block.name + prefix = f"mcp__{MCP_SERVER_NAME}__" + if tool_name.startswith(prefix): + tool_name = tool_name[len(prefix):] + logger.info("Tool call: {name}({args})", name=tool_name, args=args_str) + if turn_text: + last_text = turn_text + elif isinstance(msg, ResultMessage): + result.session_id = msg.session_id or None + result.num_turns = msg.num_turns + result.cost_usd = msg.total_cost_usd + if msg.usage: + result.input_tokens = msg.usage.get("input_tokens", 0) or 0 + result.output_tokens = msg.usage.get("output_tokens", 0) or 0 + result.cache_read_tokens = msg.usage.get("cache_read_input_tokens", 0) or 0 + result.cache_creation_tokens = msg.usage.get("cache_creation_input_tokens", 0) or 0 + if msg.result and not last_text: + last_text = msg.result result.text = last_text return result diff --git a/memclaw/backends/cursor.py b/memclaw/backends/cursor.py index 1b750a3..2a2f996 100644 --- a/memclaw/backends/cursor.py +++ b/memclaw/backends/cursor.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from collections import deque from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar @@ -29,6 +30,12 @@ # Scrub Claude credentials when switching to Cursor so they can't shadow selection. _DROP_KEYS = ["ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN"] +# INTERIM: the Cursor backend has no native conversation sessions yet, so it +# keeps the last few messages per session key in memory and folds them into +# each prompt. To be replaced by native Cursor conversation handling. +_TRANSCRIPT_MESSAGES = 10 +_TRANSCRIPT_MESSAGE_CHARS = 1000 + def _cursor_api_key(config: "MemclawConfig") -> str: return (config.cursor_api_key or os.environ.get("CURSOR_API_KEY", "")).strip() @@ -64,18 +71,31 @@ def _build_combined_prompt(*, system_prompt: str, user_message: str) -> str: ) +def _format_transcript(transcript: list[tuple[str, str]]) -> str: + """INTERIM: render recent (role, text) messages for the prompt.""" + if not transcript: + return "" + lines = ["=== RECENT CONVERSATION ==="] + for role, text in transcript: + lines.append(f"{'User' if role == 'user' else 'Assistant'}: {text}") + return "\n".join(lines) + + def _build_user_message( *, system_prompt: str, user_message: str, + context: str = "", + transcript: list[tuple[str, str]] | None = None, image_b64: str | None = None, image_media_type: str = "image/jpeg", ) -> str | Any: from cursor_sdk import SDKImage, UserMessage + parts = [_format_transcript(transcript or []), context, user_message] prompt = _build_combined_prompt( system_prompt=system_prompt, - user_message=user_message, + user_message="\n\n".join(p for p in parts if p), ) if not image_b64: return prompt @@ -117,6 +137,7 @@ class CursorAgentBackend: name: ClassVar[str] = "cursor" display_name: ClassVar[str] = "Cursor SDK" + supports_model_selection: ClassVar[bool] = False def __init__(self, config: "MemclawConfig") -> None: self.config = config @@ -126,6 +147,8 @@ def __init__(self, config: "MemclawConfig") -> None: os.environ["MEMCLAW_MEMORY_DIR"] = self._cwd self.bills_per_token = True self._mcp_server = HttpMcpServer() + # INTERIM: session key -> recent (role, text) messages. + self._transcripts: dict[str, deque[tuple[str, str]]] = {} @classmethod def is_configured(cls, config: "MemclawConfig") -> bool: @@ -215,6 +238,38 @@ async def _ensure_mcp_server(self, tool_executor: "ToolExecutor") -> None: port=self.config.mcp_http_port, ) + async def reset_session(self, session_key: str) -> None: + self._transcripts.pop(session_key, None) + + def _remember(self, session_key: str, role: str, text: str) -> None: + transcript = self._transcripts.setdefault( + session_key, deque(maxlen=_TRANSCRIPT_MESSAGES), + ) + transcript.append((role, text[:_TRANSCRIPT_MESSAGE_CHARS])) + + # Model selection isn't supported for Cursor yet (see + # supports_model_selection); CURSOR_MODEL picks the model. + + @property + def model(self) -> str: + return self._model + + @property + def effort(self) -> str | None: + return None + + async def list_models(self) -> list[Any]: + raise NotImplementedError("Model selection isn't supported for Cursor yet.") + + async def effort_levels(self) -> list[str] | None: + return None + + async def set_model(self, model_id: str) -> str: + raise NotImplementedError("Model selection isn't supported for Cursor yet.") + + async def set_effort(self, level: str) -> None: + raise NotImplementedError("Effort selection isn't supported for Cursor yet.") + async def _launch_client(self) -> Any: from cursor_sdk import AsyncClient @@ -261,8 +316,11 @@ async def run_turn( self, *, system_prompt: str, + context: str, user_message: str, tool_executor: "ToolExecutor", + session_key: str, + resume_session_id: str | None = None, image_b64: str | None = None, image_media_type: str = "image/jpeg", max_turns: int = 10, @@ -276,7 +334,9 @@ async def run_turn( message = _build_user_message( system_prompt=system_prompt, + context=context, user_message=user_message, + transcript=list(self._transcripts.get(session_key, ())), image_b64=image_b64, image_media_type=image_media_type, ) @@ -316,6 +376,8 @@ async def run_turn( if not result.text.strip(): result.text = "I couldn't generate a response." _apply_cost_fallback(result, bills_per_token=self.bills_per_token) + self._remember(session_key, "user", user_message) + self._remember(session_key, "assistant", result.text) return result except CursorAgentError as exc: logger.error("Cursor SDK turn failed: {msg}", msg=exc.message) diff --git a/memclaw/bot/handlers.py b/memclaw/bot/handlers.py index df3ef18..c043244 100644 --- a/memclaw/bot/handlers.py +++ b/memclaw/bot/handlers.py @@ -16,6 +16,7 @@ from telegram.ext import ContextTypes from ..agent import MemclawAgent +from ..backends import AgentBackend from ..config import MemclawConfig from ..reminders import ReminderScheduler from . import mask_user_id @@ -35,6 +36,89 @@ async def _typing_loop(bot, chat_id: int): pass +# Shown in Telegram's command menu (registered in cli._run_telegram). +BOT_COMMANDS = [ + ("new", "Start a fresh conversation"), + ("model", "Show or change the model"), + ("effort", "Show or change the reasoning effort"), +] + + +async def model_command_reply(backend: AgentBackend, args: list[str]) -> str: + """Reply for `/model` (list models) or `/model ` (switch).""" + if not backend.supports_model_selection: + return f"Switching models isn't supported for the {backend.display_name} backend yet." + try: + models = await backend.list_models() + except Exception as exc: + logger.warning("Could not fetch the model list: {exc}", exc=exc) + models = [] + + arg = " ".join(args).strip() + if arg: + model_id = arg + if arg.isdigit(): + n = int(arg) + if not 1 <= n <= len(models): + return f"There's no model number {n}. Send /model to see the list." + model_id = models[n - 1].id + try: + note = await backend.set_model(model_id) + except ValueError as exc: + return f"{exc} Send /model to see the available models." + lines = [f"Switched to {backend.model}."] + if note: + lines.append(note) + lines.append(f"Effort: {backend.effort or 'model default'}") + lines.append("Applies from your next message; the conversation is kept.") + return "\n".join(lines) + + lines = [ + f"Model: {backend.model}", + f"Effort: {backend.effort or 'model default'}", + "", + ] + if models: + lines.append("Available models:") + for i, m in enumerate(models, 1): + current = " (current)" if m.id == backend.model else "" + lines.append(f"{i}. {m.display_name} - {m.id}{current}") + lines += ["", "Send /model to switch."] + else: + lines.append("Couldn't load the model list right now; send /model to switch.") + return "\n".join(lines) + + +async def effort_command_reply(backend: AgentBackend, args: list[str]) -> str: + """Reply for `/effort` (show levels) or `/effort ` (set it).""" + if not backend.supports_model_selection: + return f"Changing the effort isn't supported for the {backend.display_name} backend yet." + arg = " ".join(args).strip() + if arg: + try: + await backend.set_effort(arg) + except ValueError as exc: + return str(exc) + return ( + f"Effort set to {backend.effort}.\n" + "Applies from your next message; the conversation is kept." + ) + + levels = await backend.effort_levels() + lines = [ + f"Model: {backend.model}", + f"Effort: {backend.effort or 'model default'}", + "", + ] + if levels == []: + lines.append("This model doesn't support an effort setting.") + else: + if levels: + lines.append(f"Levels this model supports: {', '.join(levels)}") + lines.append("Send /effort to change it.") + return "\n".join(lines) + + class MessageHandlers: """Routes every Telegram message through the unified Memclaw agent.""" @@ -55,7 +139,7 @@ async def _deliver_reminder(self, chat_id: str, text: str): if self._bot is None: return await self._bot.send_message(chat_id=int(chat_id), text=text) - self.agent.record_reminder_fired(text) + self.agent.record_reminder_fired(text, chat_id) def _check_user(self, user_id: int) -> bool: return user_id in self.config.allowed_user_ids_list @@ -111,7 +195,7 @@ async def _send_with_typing( await self._send_response(update, context, response_text, found_images) # ------------------------------------------------------------------ - # /start — the only command + # Commands # ------------------------------------------------------------------ async def start_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -122,9 +206,30 @@ async def start_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE "Hi! I'm your personal memory assistant powered by Memclaw.\n\n" "Just send me anything — text, photos, or voice messages.\n\n" "I'll automatically decide whether to remember it, search your " - "memories, or retrieve images. No commands needed, just talk to me." + "memories, or retrieve images. No commands needed, just talk to me.\n\n" + "/new - start a fresh conversation (your memories are kept)\n" + "/model - show or change the model\n" + "/effort - show or change the reasoning effort" ) + async def new_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE): + if not self._check_user(update.effective_user.id): + return + await self.agent.reset_conversation(str(update.effective_chat.id)) + await update.message.reply_text("Started a new conversation.") + + async def model_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE): + if not self._check_user(update.effective_user.id): + return + reply = await model_command_reply(self.agent.backend, list(context.args or [])) + await update.message.reply_text(reply) + + async def effort_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE): + if not self._check_user(update.effective_user.id): + return + reply = await effort_command_reply(self.agent.backend, list(context.args or [])) + await update.message.reply_text(reply) + # ------------------------------------------------------------------ # Message handlers — everything goes through the agent # ------------------------------------------------------------------ diff --git a/memclaw/bot/slack_handlers.py b/memclaw/bot/slack_handlers.py index c79455b..baf62f4 100644 --- a/memclaw/bot/slack_handlers.py +++ b/memclaw/bot/slack_handlers.py @@ -44,7 +44,7 @@ def __init__(self, config: MemclawConfig, openai_client: AsyncOpenAI): async def _deliver_reminder(self, chat_id: str, text: str): await self.app.client.chat_postMessage(channel=chat_id, text=text) - self.agent.record_reminder_fired(text) + self.agent.record_reminder_fired(text, chat_id) def _register_handlers(self): """Register Slack event handlers on the bolt app.""" @@ -101,7 +101,10 @@ async def _route_event(self, event: dict, say, client): image_files = [f for f in files if f.get("mimetype", "") in _IMAGE_MIMES] audio_files = [f for f in files if f.get("mimetype", "").startswith("audio/")] - if image_files: + if not files and text.lower() == "/new": + await self.agent.reset_conversation(channel) + await say(text="Started a new conversation.", thread_ts=thread_ts) + elif image_files: await self._handle_image(user, channel, thread_ts, text, image_files[0], say, client) elif audio_files: await self._handle_audio(user, channel, thread_ts, text, audio_files[0], say, client) diff --git a/memclaw/bot/whatsapp_handlers.py b/memclaw/bot/whatsapp_handlers.py index 6c0ad6c..8c4e7fc 100644 --- a/memclaw/bot/whatsapp_handlers.py +++ b/memclaw/bot/whatsapp_handlers.py @@ -50,7 +50,7 @@ async def _deliver_reminder(self, chat_id: str, text: str): user, server = chat_id, "s.whatsapp.net" jid = build_jid(user, server=server) await self.client.send_message(jid, text) - self.agent.record_reminder_fired(text) + self.agent.record_reminder_fired(text, chat_id) # ------------------------------------------------------------------ # Event registration @@ -113,7 +113,10 @@ async def _route_message(self, cli: NewAClient, ev: MessageEv): return text = msg.conversation or msg.extendedTextMessage.text - if text: + if text and text.strip().lower() == "/new": + await self.agent.reset_conversation(self._chat_id(ev)) + await cli.send_message(ev.Info.MessageSource.Chat, "Started a new conversation.") + elif text: await self._handle_text(cli, ev, text) else: logger.debug("Ignoring unsupported WhatsApp message type") diff --git a/memclaw/cli.py b/memclaw/cli.py index 394adb9..1c630b3 100644 --- a/memclaw/cli.py +++ b/memclaw/cli.py @@ -116,7 +116,8 @@ async def _interactive(config: MemclawConfig): Panel( "[bold]Memclaw[/bold] — Your Personal Memory Vault\n\n" "Type your thoughts, questions, or commands.\n" - "Type [bold]/quit[/bold] to exit.", + "Type [bold]/new[/bold] to start a fresh conversation, " + "[bold]/quit[/bold] to exit.", title="memclaw", border_style="bright_cyan", ) @@ -140,6 +141,10 @@ async def _interactive(config: MemclawConfig): break if not stripped: continue + if stripped.lower() == "/new": + await agent.reset_conversation() + console.print("[cyan]Started a new conversation.[/cyan]\n") + continue try: with console.status("[cyan]Thinking...[/cyan]"): @@ -178,10 +183,11 @@ def _setup_bot_logging(log_file: Path) -> None: def _run_telegram(config: MemclawConfig) -> None: from loguru import logger from openai import AsyncOpenAI - from telegram.error import NetworkError, TimedOut + from telegram import BotCommand + from telegram.error import NetworkError, TelegramError, TimedOut from telegram.ext import Application, CommandHandler, MessageHandler, filters - from .bot.handlers import MessageHandlers + from .bot.handlers import BOT_COMMANDS, MessageHandlers if not config.telegram_bot_token: console.print("[red]Error:[/red] TELEGRAM_BOT_TOKEN is not set.") @@ -201,6 +207,13 @@ async def post_init(application: Application) -> None: handlers.attach_bot(application.bot) handlers.scheduler.start() + try: + await application.bot.set_my_commands( + [BotCommand(name, description) for name, description in BOT_COMMANDS] + ) + except TelegramError as exc: + logger.warning(f"Could not register the bot command menu: {exc}") + logger.info("Memclaw bot initialized") async def post_shutdown(application: Application) -> None: @@ -220,6 +233,15 @@ async def post_shutdown(application: Application) -> None: async def _start(update, context): await context.bot_data["handlers"].start_command(update, context) + async def _new(update, context): + await context.bot_data["handlers"].new_command(update, context) + + async def _model(update, context): + await context.bot_data["handlers"].model_command(update, context) + + async def _effort(update, context): + await context.bot_data["handlers"].effort_command(update, context) + async def _text(update, context): await context.bot_data["handlers"].handle_text(update, context) @@ -238,6 +260,9 @@ async def _on_error(update, context): app.add_error_handler(_on_error) app.add_handler(CommandHandler("start", _start)) + app.add_handler(CommandHandler("new", _new)) + app.add_handler(CommandHandler("model", _model)) + app.add_handler(CommandHandler("effort", _effort)) app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, _text)) app.add_handler(MessageHandler(filters.PHOTO, _photo)) app.add_handler(MessageHandler(filters.VOICE, _voice)) diff --git a/memclaw/config.py b/memclaw/config.py index 0ae9acc..97ca69d 100644 --- a/memclaw/config.py +++ b/memclaw/config.py @@ -49,10 +49,6 @@ class MemclawConfig: cursor_model: str = "" mcp_http_port: int = 17373 - # Conversation continuity - conversation_history_limit: int = 10 - conversation_history_window_minutes: int = 60 - # Memory consolidation consolidation_threshold: int = 7 diff --git a/memclaw/defaults/AGENTS.md b/memclaw/defaults/AGENTS.md index 140f9d6..aeb054c 100644 --- a/memclaw/defaults/AGENTS.md +++ b/memclaw/defaults/AGENTS.md @@ -9,7 +9,7 @@ You are Memclaw, a personal memory assistant. You help users store and retrieve 3. **Images**: When you see an image with a media_ref (from a messaging channel or a local path), describe what you see in detail and save using image_save. Pass the media_ref verbatim. 4. **Image retrieval**: When the user asks to find an image — use image_search. The image will be sent automatically. 5. **Links**: URLs in the user's message arrive pre-fetched and summarised as `[Link summary] : `. If worth keeping, save the URL + summary with memory_save. Tell users you can do this when asked. -6. **Reminders**: When the user asks to be reminded — use reminder_create. One-shot ("remind me tomorrow at 9am to call Alex") sets `fire_at` to a local ISO 8601 datetime. Recurring ("every 5 hours drink water") sets `interval_seconds`. Use the "Current local time" in the system prompt as the reference for converting natural language. Use reminder_list to show pending ones, reminder_cancel to remove by id. Do NOT also memory_save reminders — they live in their own table. +6. **Reminders**: When the user asks to be reminded — use reminder_create. One-shot ("remind me tomorrow at 9am to call Alex") sets `fire_at` to a local ISO 8601 datetime. Recurring ("every 5 hours drink water") sets `interval_seconds`. Use the "Current local time" in the block at the top of each message as the reference for converting natural language. Use reminder_list to show pending ones, reminder_cancel to remove by id. Do NOT also memory_save reminders — they live in their own table. 7. **Conversation**: Sometimes the user just wants to chat. Respond naturally. If they mention something worth remembering, save it too. ## Storage guidelines diff --git a/memclaw/sessions.py b/memclaw/sessions.py new file mode 100644 index 0000000..a5f0f1f --- /dev/null +++ b/memclaw/sessions.py @@ -0,0 +1,65 @@ +"""Persisted per-chat conversation sessions. + +Maps a chat's session key to the backend session id the conversation lives +in, so a restarted Memclaw resumes it. Each entry also records: + +- ``fingerprint`` — hash of the code-level instructions (system-prompt + template + backend) the session started with. Claude Code keeps a + session's original system prompt on resume, so a session started under + different instructions is not resumed. +- ``agents_hash`` / ``memory_hash`` — hashes of the AGENTS.md and MEMORY.md + content the session has seen, so a later change is sent to it once. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +from loguru import logger + + +class SessionStore: + """``{session_key: {session_id, fingerprint, agents_hash, memory_hash}}`` + in a JSON file, written atomically.""" + + def __init__(self, path: Path) -> None: + self.path = path + self._data: dict[str, dict[str, str]] = {} + try: + raw = json.loads(path.read_text()) + except FileNotFoundError: + raw = {} + except (OSError, ValueError) as exc: + logger.warning("Ignoring unreadable session store {p}: {exc}", p=path, exc=exc) + raw = {} + if isinstance(raw, dict): + self._data = {k: v for k, v in raw.items() if isinstance(v, dict)} + + def get(self, key: str, fingerprint: str) -> dict[str, str] | None: + """The entry to resume, or None. An entry from other instructions is + dropped.""" + entry = self._data.get(key) + if entry is None or not entry.get("session_id"): + return None + if entry.get("fingerprint") != fingerprint: + logger.info("Instructions changed; not resuming the old session for {k}", k=key) + self.drop(key) + return None + return dict(entry) + + def set(self, key: str, entry: dict[str, str]) -> None: + if self._data.get(key) == entry: + return + self._data[key] = dict(entry) + self._save() + + def drop(self, key: str) -> None: + if self._data.pop(key, None) is not None: + self._save() + + def _save(self) -> None: + self.path.parent.mkdir(parents=True, exist_ok=True) + tmp = self.path.with_suffix(".tmp") + tmp.write_text(json.dumps(self._data, indent=2)) + tmp.replace(self.path) diff --git a/memclaw/setup.py b/memclaw/setup.py index d96dbf4..d1d84c0 100644 --- a/memclaw/setup.py +++ b/memclaw/setup.py @@ -206,6 +206,25 @@ def _load_existing() -> dict[str, str]: return values +def update_env_file(values: dict[str, str], drop_keys: list[str] | tuple[str, ...] = ()) -> None: + """Merge *values* into ~/.memclaw/.env and remove *drop_keys*. + + Keys are written in their existing order, new ones appended; a key with + an empty value is removed, matching how the wizard writes the file. + """ + merged = _load_existing() + merged.update(values) + for key in drop_keys: + merged.pop(key, None) + _write_env(merged) + + +def _write_env(values: dict[str, str]) -> None: + ENV_FILE.parent.mkdir(parents=True, exist_ok=True) + lines = [f"{k}={v}" for k, v in values.items() if v] + ENV_FILE.write_text("\n".join(lines) + "\n") + + def needs_setup() -> bool: """Return True if first-run setup is needed.""" return not ENV_FILE.exists() @@ -298,9 +317,7 @@ def _is_required(required: bool, key_platform: str | None) -> bool: raise SystemExit(1) # Write to ~/.memclaw/.env - ENV_FILE.parent.mkdir(parents=True, exist_ok=True) - lines = [f"{k}={v}" for k, v in values.items() if v] - ENV_FILE.write_text("\n".join(lines) + "\n") + _write_env(values) # Sync the current process env with the freshly-written .env. Without this, # MemclawConfig.__post_init__ would still see stale values (e.g. an OAuth diff --git a/pyproject.toml b/pyproject.toml index a96968e..fbeecc7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "claude-agent-sdk>=0.2.139,<0.3", + "claude-agent-sdk>=0.2.163,<0.3", "mcp>=1.23.0,<2", "openai>=1.0.0", "numpy>=1.24.0", diff --git a/tests/backends/test_claude.py b/tests/backends/test_claude.py index d1ea1bc..0381b63 100644 --- a/tests/backends/test_claude.py +++ b/tests/backends/test_claude.py @@ -20,7 +20,7 @@ pytestmark = pytest.mark.usefixtures("isolate_claude_env") -def _mock_sdk_client(text: str): +def _mock_sdk_client(text: str, *, session_id: str = "test", costs=None): """Build a fake ClaudeSDKClient context that yields one AssistantMessage followed by a ResultMessage. Returns (ctx_factory, client_mock). """ @@ -28,7 +28,7 @@ def _mock_sdk_client(text: str): assistant = AssistantMessage( content=[TextBlock(text=text)], - model="claude-sonnet-4-6", + model="claude-sonnet-5-5", ) result = ResultMessage( subtype="result", @@ -36,7 +36,7 @@ def _mock_sdk_client(text: str): duration_api_ms=8, is_error=False, num_turns=1, - session_id="test", + session_id=session_id, total_cost_usd=None, usage={"input_tokens": 100, "output_tokens": 50, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0}, @@ -46,9 +46,13 @@ def _mock_sdk_client(text: str): client = MagicMock() client.query = AsyncMock() + # *costs*: the running session total the CLI reports, one per query. + cost_iter = iter(costs or ()) + def _receive_factory(): async def _gen(): yield assistant + result.total_cost_usd = next(cost_iter, None) yield result return _gen() @@ -63,6 +67,15 @@ def _ctx_factory(*args, **kwargs): return _ctx_factory, client +def _model_info(model_id: str, *, effort_levels: list[str] | None = None): + return ModelInfo( + id=model_id, + display_name=model_id, + created_at="2026-07-24T00:00:00Z", + effort_levels=effort_levels or [], + ) + + # ──────────────────────────────────────────────────────────────────── # Auth mode + env scrubbing # ──────────────────────────────────────────────────────────────────── @@ -160,31 +173,61 @@ async def test_returns_text(self, claude_config): assert client.query.await_args.args[0] == "hi" +def _executor(cfg): + from memclaw.tools import ToolExecutor + + return ToolExecutor( + config=cfg, store=MagicMock(), index=MagicMock(), + search=MagicMock(), found_images=[], platform="test", + ) + + +class _FakeClients: + """Stands in for ClaudeSDKClient: records every client built (with its + options) and serves one canned reply per query.""" + + def __init__(self, text: str = "done", *, session_id: str = "sess-1", + fail_resume: bool = False, costs=None): + self.costs = costs + self.text = text + self.session_id = session_id + self.fail_resume = fail_resume + self.clients: list[MagicMock] = [] + self.options: list = [] + + def __call__(self, options=None, *args, **kwargs): + _factory, client = _mock_sdk_client( + self.text, session_id=self.session_id, costs=self.costs, + ) + resume = options.resume if options is not None else None + client.connect = AsyncMock( + side_effect=RuntimeError("no such session") if self.fail_resume and resume else None, + ) + client.disconnect = AsyncMock() + self.clients.append(client) + self.options.append(options) + return client + + +async def _turn(backend, executor, *, key="telegram:1", resume=None, + context="now", message="hello", **kwargs): + return await backend.run_turn( + system_prompt="sys", context=context, user_message=message, + tool_executor=executor, session_key=key, resume_session_id=resume, + **kwargs, + ) + + class TestRunTurn: @pytest.mark.asyncio async def test_returns_turn_result(self, claude_config): - from memclaw.tools import ToolExecutor - backend = ClaudeAgentBackend(claude_config(api_key="x")) - cfg = backend.config - executor = ToolExecutor( - config=cfg, - store=MagicMock(), - index=MagicMock(), - search=MagicMock(), - found_images=[], - platform="test", - ) - - ctx_factory, _client = _mock_sdk_client("done") - with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=ctx_factory): - result = await backend.run_turn( - system_prompt="sys", - user_message="hello", - tool_executor=executor, - ) + fake = _FakeClients("done", session_id="sess-1") + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + result = await _turn(backend, _executor(backend.config)) assert result.text == "done" + assert result.session_id == "sess-1" assert result.input_tokens == 100 assert result.output_tokens == 50 # bills_per_token is True (api_key) and the mock returns no cost, @@ -194,31 +237,319 @@ async def test_returns_turn_result(self, claude_config): @pytest.mark.asyncio async def test_fallback_cost_uses_the_configured_models_price(self, claude_config): - from memclaw.tools import ToolExecutor - cfg = claude_config(api_key="x") cfg.claude_model = "claude-opus-5" backend = ClaudeAgentBackend(cfg) - executor = ToolExecutor( - config=cfg, store=MagicMock(), index=MagicMock(), - search=MagicMock(), found_images=[], platform="test", - ) - - ctx_factory, _client = _mock_sdk_client("done") - with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=ctx_factory): - result = await backend.run_turn( - system_prompt="sys", user_message="hello", tool_executor=executor, - ) + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=_FakeClients()): + result = await _turn(backend, _executor(cfg)) # 100 input tokens at $5/M + 50 output tokens at $25/M (Opus 5). assert result.cost_usd == pytest.approx((100 * 5 + 50 * 25) / 1_000_000) + @pytest.mark.asyncio + async def test_cost_is_per_turn_not_the_session_running_total(self, claude_config): + backend = ClaudeAgentBackend(claude_config(api_key="x")) + fake = _FakeClients(costs=[0.01, 0.025]) + executor = _executor(backend.config) + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + first = await _turn(backend, executor) + second = await _turn(backend, executor, resume="sess-1") + + assert first.cost_usd == pytest.approx(0.01) + assert second.cost_usd == pytest.approx(0.015) + + @pytest.mark.asyncio + async def test_resumed_session_without_baseline_estimates_cost(self, claude_config): + # After a restart the running total includes earlier turns, so the + # token-based estimate is used instead. + backend = ClaudeAgentBackend(claude_config(api_key="x")) + fake = _FakeClients(costs=[5.0]) + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + result = await _turn(backend, _executor(backend.config), resume="sess-1") + + input_per_m, output_per_m = claude_backend._prices_per_m(backend.model) + assert result.cost_usd == pytest.approx((100 * input_per_m + 50 * output_per_m) / 1_000_000) + + @pytest.mark.asyncio + async def test_session_options_are_locked_down(self, claude_config): + backend = ClaudeAgentBackend(claude_config(oauth="x")) + fake = _FakeClients() + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + await _turn(backend, _executor(backend.config), max_turns=7) + + options = fake.options[0] + assert options.setting_sources == [] + assert options.strict_mcp_config is True + assert options.tools == [] + assert options.permission_mode == "dontAsk" + assert options.verbatim_prompts is True + assert options.allowed_tools == claude_backend._ALLOWED_TOOLS + assert all(t.startswith("mcp__memclaw__") for t in options.allowed_tools) + assert options.disallowed_tools == [] + assert options.system_prompt == "sys" + assert options.cwd == str(backend.config.memory_dir) + assert options.max_turns == 7 + assert options.resume is None + + @pytest.mark.asyncio + async def test_context_block_is_prefixed_to_the_message(self, claude_config): + backend = ClaudeAgentBackend(claude_config(oauth="x")) + fake = _FakeClients() + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + await _turn(backend, _executor(backend.config), + context="ctx", message="hi there") + + prompt = fake.clients[0].query.await_args.args[0] + assert prompt == "ctx\n\nhi there" + + @pytest.mark.asyncio + async def test_image_turn_streams_context_in_the_text_part(self, claude_config): + backend = ClaudeAgentBackend(claude_config(oauth="x")) + fake = _FakeClients() + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + await _turn(backend, _executor(backend.config), + context="ctx", message="a photo", + image_b64="aGVsbG8=", image_media_type="image/png") + + stream = fake.clients[0].query.await_args.args[0] + messages = [m async for m in stream] + content = messages[0]["message"]["content"] + assert content[0]["source"] == { + "type": "base64", "media_type": "image/png", "data": "aGVsbG8=", + } + assert content[1] == {"type": "text", "text": "ctx\n\na photo"} + + +class TestPersistentSessions: + @pytest.mark.asyncio + async def test_client_is_reused_across_turns(self, claude_config): + backend = ClaudeAgentBackend(claude_config(oauth="x")) + fake = _FakeClients() + executor = _executor(backend.config) + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + await _turn(backend, executor, message="one") + await _turn(backend, executor, message="two", resume="sess-1") + + assert len(fake.clients) == 1 + client = fake.clients[0] + client.connect.assert_awaited_once() + assert client.query.await_count == 2 + client.disconnect.assert_not_awaited() + + @pytest.mark.asyncio + async def test_each_chat_gets_its_own_client(self, claude_config): + backend = ClaudeAgentBackend(claude_config(oauth="x")) + fake = _FakeClients() + executor = _executor(backend.config) + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + await _turn(backend, executor, key="telegram:1") + await _turn(backend, executor, key="telegram:2") + assert len(fake.clients) == 2 + + @pytest.mark.asyncio + async def test_resume_id_reaches_the_cli(self, claude_config): + backend = ClaudeAgentBackend(claude_config(oauth="x")) + fake = _FakeClients() + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + await _turn(backend, _executor(backend.config), resume="old-session") + assert fake.options[0].resume == "old-session" + + @pytest.mark.asyncio + async def test_failed_resume_falls_back_to_a_fresh_session(self, claude_config): + backend = ClaudeAgentBackend(claude_config(oauth="x")) + fake = _FakeClients("fresh", session_id="new-session", fail_resume=True) + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + result = await _turn(backend, _executor(backend.config), resume="gone") + + assert [o.resume for o in fake.options] == ["gone", None] + assert result.text == "fresh" + assert result.session_id == "new-session" + + @pytest.mark.asyncio + async def test_reset_session_closes_the_client(self, claude_config): + backend = ClaudeAgentBackend(claude_config(oauth="x")) + fake = _FakeClients() + executor = _executor(backend.config) + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + await _turn(backend, executor) + await backend.reset_session("telegram:1") + fake.clients[0].disconnect.assert_awaited_once() + await _turn(backend, executor) + + assert len(fake.clients) == 2 + assert fake.options[1].resume is None + + @pytest.mark.asyncio + async def test_reset_of_an_unknown_session_is_a_no_op(self, claude_config): + backend = ClaudeAgentBackend(claude_config(oauth="x")) + await backend.reset_session("nope") + + @pytest.mark.asyncio + async def test_failed_turn_drops_the_client_so_the_next_one_resumes(self, claude_config): + backend = ClaudeAgentBackend(claude_config(oauth="x")) + fake = _FakeClients() + executor = _executor(backend.config) + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + await _turn(backend, executor) + fake.clients[0].query = AsyncMock(side_effect=RuntimeError("CLI died")) + with pytest.raises(RuntimeError): + await _turn(backend, executor, resume="sess-1") + fake.clients[0].disconnect.assert_awaited_once() + await _turn(backend, executor, resume="sess-1") + + assert len(fake.clients) == 2 + assert fake.options[1].resume == "sess-1" + + @pytest.mark.asyncio + async def test_model_change_reconnects_resuming_the_session( + self, claude_config, tmp_path, monkeypatch, + ): + monkeypatch.setattr("memclaw.setup.ENV_FILE", tmp_path / ".env") + backend = ClaudeAgentBackend(claude_config(oauth="x")) + backend._models = [_model_info("claude-opus-5", effort_levels=["low", "medium"])] + fake = _FakeClients() + executor = _executor(backend.config) + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + await _turn(backend, executor) + await backend.set_model("claude-opus-5") + # The live client is only replaced on the next turn. + fake.clients[0].disconnect.assert_not_awaited() + await _turn(backend, executor, resume="sess-1") + + fake.clients[0].disconnect.assert_awaited_once() + assert len(fake.clients) == 2 + assert fake.options[1].resume == "sess-1" + assert fake.options[1].model == "claude-opus-5" + + @pytest.mark.asyncio + async def test_shutdown_closes_every_client(self, claude_config): + backend = ClaudeAgentBackend(claude_config(oauth="x")) + fake = _FakeClients() + executor = _executor(backend.config) + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + await _turn(backend, executor, key="a") + await _turn(backend, executor, key="b") + await backend.on_agent_shutdown() + for client in fake.clients: + client.disconnect.assert_awaited_once() + + +class TestModelSelection: + @pytest.fixture + def env_file(self, tmp_path, monkeypatch): + path = tmp_path / ".env" + path.write_text("OPENAI_API_KEY=k\nCLAUDE_EFFORT=high\n") + monkeypatch.setattr("memclaw.setup.ENV_FILE", path) + return path + + MODELS = [ + _model_info("claude-opus-5", effort_levels=["low", "medium", "high", "max"]), + _model_info("claude-sonnet-5-5", effort_levels=["low", "medium"]), + _model_info("claude-haiku-4-5"), + ] + + def _backend(self, claude_config, effort="high"): + cfg = claude_config(oauth="x") + cfg.claude_model = "claude-opus-5" + cfg.claude_effort = effort + backend = ClaudeAgentBackend(cfg) + return backend + + @pytest.mark.asyncio + async def test_list_models_uses_the_configured_credential(self, claude_config): + backend = ClaudeAgentBackend(claude_config(oauth="tok")) + fetch = AsyncMock(return_value=self.MODELS) + with patch("memclaw.backends.claude.fetch_models", fetch): + assert await backend.list_models() == self.MODELS + fetch.assert_awaited_once_with("subscription", "tok") + + @pytest.mark.asyncio + async def test_set_model_persists_and_updates_config(self, claude_config, env_file): + backend = self._backend(claude_config) + backend._models = list(self.MODELS) + note = await backend.set_model("claude-opus-5") + assert note == "" + assert backend.model == "claude-opus-5" + assert backend.effort == "high" + assert backend.config.claude_model == "claude-opus-5" + text = env_file.read_text() + assert "CLAUDE_MODEL=claude-opus-5" in text + assert "CLAUDE_EFFORT=high" in text + assert "OPENAI_API_KEY=k" in text + + @pytest.mark.asyncio + async def test_unsupported_effort_steps_down(self, claude_config, env_file): + backend = self._backend(claude_config) + backend._models = list(self.MODELS) + note = await backend.set_model("claude-sonnet-5-5") + assert backend.effort == "medium" + assert "'high'" in note and "'medium'" in note + assert "CLAUDE_EFFORT=medium" in env_file.read_text() + + @pytest.mark.asyncio + async def test_model_without_effort_drops_it(self, claude_config, env_file): + backend = self._backend(claude_config) + backend._models = list(self.MODELS) + note = await backend.set_model("claude-haiku-4-5") + assert backend.effort is None + assert "effort" in note + text = env_file.read_text() + assert "CLAUDE_MODEL=claude-haiku-4-5" in text + assert "CLAUDE_EFFORT" not in text + + @pytest.mark.asyncio + async def test_unknown_model_is_rejected(self, claude_config, env_file): + backend = self._backend(claude_config) + backend._models = list(self.MODELS) + with pytest.raises(ValueError, match="Unknown model"): + await backend.set_model("claude-nope") + assert backend.model == "claude-opus-5" + + @pytest.mark.asyncio + async def test_any_id_is_accepted_when_the_list_is_unavailable( + self, claude_config, env_file, + ): + backend = self._backend(claude_config) + with patch("memclaw.backends.claude.fetch_models", + AsyncMock(side_effect=httpx.ConnectError("offline"))): + await backend.set_model("claude-future-9") + assert backend.model == "claude-future-9" + assert backend.effort == "high" + + @pytest.mark.asyncio + async def test_set_effort(self, claude_config, env_file): + backend = self._backend(claude_config) + backend._models = list(self.MODELS) + await backend.set_effort(" MAX ") + assert backend.effort == "max" + assert "CLAUDE_EFFORT=max" in env_file.read_text() + + @pytest.mark.asyncio + async def test_set_effort_rejects_unsupported_levels(self, claude_config, env_file): + backend = self._backend(claude_config) + backend._models = list(self.MODELS) + with pytest.raises(ValueError, match="Unknown effort"): + await backend.set_effort("ultra") + with pytest.raises(ValueError, match="doesn't support 'xhigh'"): + await backend.set_effort("xhigh") + assert backend.effort == "high" + + @pytest.mark.asyncio + async def test_set_effort_on_a_model_without_effort(self, claude_config, env_file): + backend = self._backend(claude_config, effort="") + backend._model = "claude-haiku-4-5" + backend._models = list(self.MODELS) + assert await backend.effort_levels() == [] + with pytest.raises(ValueError, match="doesn't support an effort"): + await backend.set_effort("low") + class TestPrices: @pytest.mark.parametrize("model, prices", [ ("claude-opus-5", (5.0, 25.0)), ("claude-opus-5-5", (4.0, 20.0)), # not mistaken for claude-opus-5 ("claude-sonnet-4-6", (3.0, 15.0)), + ("claude-sonnet-5-5", (2.0, 10.0)), ("claude-haiku-4-5-20251001", (1.0, 5.0)), # dated snapshot ]) def test_known_models(self, model, prices): @@ -252,8 +583,14 @@ def test_blank_value_falls_back(self, claude_config): class TestResolveEffort: - def test_unset_gives_none(self, claude_config): + def test_unset_gives_the_default_on_the_default_model(self, claude_config): + cfg = claude_config(oauth="x") + assert claude_backend._resolve_effort(cfg) == claude_backend._DEFAULT_EFFORT == "medium" + + def test_unset_gives_none_on_another_model(self, claude_config): + """Another model may not take an effort at all, so leave it to the model.""" cfg = claude_config(oauth="x") + cfg.claude_model = "claude-haiku-4-5" assert claude_backend._resolve_effort(cfg) is None def test_known_level_passes_through(self, claude_config): @@ -269,6 +606,7 @@ def test_case_and_padding_are_normalised(self, claude_config): def test_level_the_sdk_does_not_know_is_ignored(self, claude_config): """A typo in ~/.memclaw/.env must not reach the CLI as an argument.""" cfg = claude_config(oauth="x") + cfg.claude_model = "claude-opus-5" cfg.claude_effort = "ultra" assert claude_backend._resolve_effort(cfg) is None @@ -284,7 +622,7 @@ def test_configured_values_are_shown(self, claude_config): def test_defaults_are_shown_when_unset(self, claude_config): rows = ClaudeAgentBackend.status_rows(claude_config(api_key="k")) - assert rows == [("Model", claude_backend._DEFAULT_MODEL), ("Effort", "default")] + assert rows == [("Model", "claude-sonnet-5-5"), ("Effort", "medium")] class TestOptionsCarryModelAndEffort: @@ -322,23 +660,34 @@ def _capture(options, *args, **kwargs): with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=_capture): await backend.run_one_shot(system_prompt="s", user_message="u") - assert seen["options"].model == claude_backend._DEFAULT_MODEL - assert seen["options"].effort is None + assert seen["options"].model == "claude-sonnet-5-5" + assert seen["options"].effort == "medium" + + @pytest.mark.asyncio + async def test_one_shot_uses_the_fast_options(self, claude_config): + backend = ClaudeAgentBackend(claude_config(oauth="x")) + + ctx_factory, _client = _mock_sdk_client("ok") + seen = {} + + def _capture(options, *args, **kwargs): + seen["options"] = options + return ctx_factory() + + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=_capture): + await backend.run_one_shot(system_prompt="s", user_message="u") + + options = seen["options"] + assert options.setting_sources == [] + assert options.strict_mcp_config is True + assert options.tools == [] + assert options.max_turns == 1 # ──────────────────────────────────────────────────────────────────── # Wizard: model + effort questions # ──────────────────────────────────────────────────────────────────── -def _model_info(model_id: str, *, effort_levels: list[str] | None = None): - return ModelInfo( - id=model_id, - display_name=model_id, - created_at="2026-07-24T00:00:00Z", - effort_levels=effort_levels or [], - ) - - def _ask(models, *, existing=None, answers=("1",), fetch_error: Exception | None = None): """Run `_ask_model_and_effort` with the fetch and the prompts mocked. @@ -383,13 +732,13 @@ def test_effort_question_is_skipped_without_support(self): assert drops == ["CLAUDE_EFFORT"] assert len(prompts) == 1 - def test_effort_defaults_to_high(self): + def test_effort_defaults_to_medium(self): models = [_model_info("claude-opus-5", effort_levels=["low", "medium", "high"])] - _, _, prompts, _ = _ask(models, answers=("1", "3")) - assert prompts[1]["default"] == "3" + _, _, prompts, _ = _ask(models, answers=("1", "2")) + assert prompts[1]["default"] == "2" - def test_effort_default_falls_back_when_high_is_unsupported(self): - models = [_model_info("claude-opus-5", effort_levels=["low", "medium"])] + def test_effort_default_falls_back_when_medium_is_unsupported(self): + models = [_model_info("claude-opus-5", effort_levels=["low", "high"])] _, _, prompts, _ = _ask(models, answers=("1", "1")) assert prompts[1]["default"] == "1" diff --git a/tests/backends/test_cursor.py b/tests/backends/test_cursor.py index 7b1a29a..4744e5d 100644 --- a/tests/backends/test_cursor.py +++ b/tests/backends/test_cursor.py @@ -323,8 +323,10 @@ async def test_run_turn(self, tmp_path): result = await backend.run_turn( system_prompt="System", + context="ctx", user_message="User", tool_executor=MagicMock(), + session_key="cli", max_turns=5, ) @@ -383,8 +385,10 @@ async def _send(message, options): result = await backend.run_turn( system_prompt="System", + context="ctx", user_message="User", tool_executor=MagicMock(), + session_key="cli", max_turns=5, ) @@ -406,10 +410,68 @@ async def test_run_turn_without_mcp_server_raises(self, tmp_path): with pytest.raises(RuntimeError, match="MCP server failed to start"): await backend.run_turn( system_prompt="System", + context="", user_message="User", tool_executor=MagicMock(), + session_key="cli", ) + @pytest.mark.asyncio + async def test_interim_transcript_is_folded_into_later_prompts(self, tmp_path): + cfg = _make_config(tmp_path, cursor_api_key="crsr_test_key") + backend = CursorAgentBackend(cfg) + + sent: list[str] = [] + + def _make_run(reply): + run = AsyncMock() + run.events = MagicMock(return_value=_empty_events()) + run.wait = AsyncMock(return_value=SimpleNamespace(result=reply, num_turns=1)) + return run + + replies = iter(["First reply", "Second reply", "Third reply"]) + + async def _send(message, options): + sent.append(message) + return _make_run(next(replies)) + + mock_agent = AsyncMock() + mock_agent.send = AsyncMock(side_effect=_send) + mock_agent.close = AsyncMock() + mock_client = AsyncMock() + mock_client.agents.create = AsyncMock(return_value=mock_agent) + mock_mcp = HttpMcpServerConfig(url="http://127.0.0.1:8765/mcp", type="http") + + async def _turn(key, text): + return await backend.run_turn( + system_prompt="System", context="ctx", + user_message=text, tool_executor=MagicMock(), session_key=key, + ) + + with patch("cursor_sdk.AsyncClient") as mock_client_cls, patch.object( + backend, "_ensure_mcp_server", new_callable=AsyncMock + ), patch.object( + HttpMcpServer, "config", new_callable=PropertyMock, return_value=mock_mcp + ): + mock_client_cls.launch_bridge = AsyncMock(return_value=mock_client) + await _turn("cli", "My name is Ana") + await _turn("cli", "What's my name?") + await backend.reset_session("cli") + await _turn("cli", "Fresh start") + + assert "RECENT CONVERSATION" not in sent[0] + assert "ctx" in sent[0] + assert "User: My name is Ana" in sent[1] + assert "Assistant: First reply" in sent[1] + # The context block itself isn't kept in the transcript. + assert sent[1].count("ctx") == 1 + assert "RECENT CONVERSATION" not in sent[2] + + def test_model_selection_is_not_supported_yet(self, tmp_path): + backend = CursorAgentBackend(_make_config(tmp_path, cursor_api_key="k")) + assert backend.supports_model_selection is False + assert backend.model == "composer-2.5" + @pytest.mark.asyncio async def test_agent_start_starts_mcp_server(self, tmp_path, monkeypatch): monkeypatch.setenv("AGENT_BACKEND", "cursor") diff --git a/tests/test_agent.py b/tests/test_agent.py index 354bd6e..2b50443 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -1,4 +1,4 @@ -"""Tests for MemclawAgent — history, consolidation, context, sync, fs guardrail. +"""Tests for MemclawAgent — sessions, consolidation, context, sync, fs guardrail. These tests run against a backend-agnostic `FakeBackend` so they exercise the orchestration in `MemclawAgent` without depending on any SDK. Per-backend @@ -7,8 +7,9 @@ """ from __future__ import annotations +import asyncio import json -from datetime import date, datetime, timedelta +from datetime import date from pathlib import Path from typing import Any from unittest.mock import AsyncMock @@ -49,6 +50,10 @@ class FakeBackend: turn_response: TurnResult = field(default_factory=lambda: TurnResult(text="")) one_shot_calls: list[dict[str, Any]] = field(default_factory=list) turn_calls: list[dict[str, Any]] = field(default_factory=list) + reset_calls: list[str] = field(default_factory=list) + # Session id each turn reports; None mimics a backend without sessions. + session_id: str | None = "sess-1" + turn_error: Exception | None = None async def on_agent_start(self, tool_executor) -> None: pass @@ -62,7 +67,14 @@ async def run_one_shot(self, *, system_prompt: str, user_message: str) -> str: async def run_turn(self, **kwargs) -> TurnResult: self.turn_calls.append(kwargs) - return self.turn_response + if self.turn_error is not None: + raise self.turn_error + result = TurnResult(**vars(self.turn_response)) + result.session_id = self.session_id + return result + + async def reset_session(self, session_key: str) -> None: + self.reset_calls.append(session_key) @pytest.fixture @@ -70,176 +82,329 @@ def fake_backend() -> FakeBackend: return FakeBackend() -def _make_agent(cfg: MemclawConfig, backend: FakeBackend): +def _make_agent(cfg: MemclawConfig, backend: FakeBackend, platform: str | None = None): from memclaw.agent import MemclawAgent - return MemclawAgent(cfg, backend=backend) + return MemclawAgent(cfg, platform=platform, backend=backend) + + +def _chat_agent(cfg: MemclawConfig, backend: FakeBackend, platform: str | None = None): + """An agent ready for handle(): no search, no consolidation.""" + agent = _make_agent(cfg, backend, platform) + agent.search.search = AsyncMock(return_value=[]) + agent._maybe_consolidate = AsyncMock(return_value=False) + return agent # ──────────────────────────────────────────────────────────────────── -# Spec #1: Conversation History +# Conversation sessions # ──────────────────────────────────────────────────────────────────── -class TestConversationHistory: - def test_history_initialized_empty(self, cfg: MemclawConfig, fake_backend: FakeBackend): - agent = _make_agent(cfg, fake_backend) - assert agent._history == [] +class TestSessions: + @pytest.mark.asyncio + async def test_first_turn_starts_a_session(self, cfg: MemclawConfig, fake_backend: FakeBackend): + fake_backend.turn_response = TurnResult(text="Hello! I'm Memclaw.") + agent = _chat_agent(cfg, fake_backend) + + reply, _ = await agent.handle("Hello") + + assert reply == "Hello! I'm Memclaw." + call = fake_backend.turn_calls[0] + assert call["user_message"] == "Hello" + assert call["session_key"] == "cli" + assert call["resume_session_id"] is None + stored = json.loads((cfg.memory_dir / "sessions.json").read_text()) + assert stored["cli"]["session_id"] == "sess-1" agent.close() @pytest.mark.asyncio - async def test_history_appended_after_handle(self, cfg: MemclawConfig, fake_backend: FakeBackend): - """handle() should append user + assistant messages to _history.""" - fake_backend.turn_response = TurnResult(text="Hello! I'm Memclaw.") - agent = _make_agent(cfg, fake_backend) + async def test_session_key_is_per_chat_and_platform( + self, cfg: MemclawConfig, fake_backend: FakeBackend, + ): + agent = _chat_agent(cfg, fake_backend, platform="telegram") + await agent.handle("hi", chat_id="42") + assert fake_backend.turn_calls[0]["session_key"] == "telegram:42" + agent.close() - agent.build_context = AsyncMock(return_value="No memories found yet.") - agent._maybe_consolidate = AsyncMock(return_value=False) + @pytest.mark.asyncio + async def test_next_turn_and_restart_resume_the_session( + self, cfg: MemclawConfig, fake_backend: FakeBackend, + ): + agent = _chat_agent(cfg, fake_backend) + await agent.handle("one") + await agent.handle("two") + assert fake_backend.turn_calls[1]["resume_session_id"] == "sess-1" + agent.close() - await agent.handle("Hello") + restarted = _chat_agent(cfg, fake_backend) + await restarted.handle("three") + assert fake_backend.turn_calls[2]["resume_session_id"] == "sess-1" + restarted.close() - assert len(agent._history) == 2 - assert agent._history[0]["role"] == "user" - assert agent._history[0]["content"] == "Hello" - assert agent._history[1]["role"] == "assistant" - assert agent._history[1]["content"] == "Hello! I'm Memclaw." - assert "timestamp" in agent._history[0] - # Backend was called exactly once with the user message. - assert len(fake_backend.turn_calls) == 1 - assert fake_backend.turn_calls[0]["user_message"] == "Hello" + @pytest.mark.asyncio + async def test_new_session_id_after_a_failed_resume_is_stored( + self, cfg: MemclawConfig, fake_backend: FakeBackend, + ): + agent = _chat_agent(cfg, fake_backend) + await agent.handle("one") + fake_backend.session_id = "sess-2" # backend couldn't resume, started fresh + await agent.handle("two") + await agent.handle("three") + assert fake_backend.turn_calls[2]["resume_session_id"] == "sess-2" agent.close() - def test_history_trimming(self, cfg: MemclawConfig, fake_backend: FakeBackend): - """Floor case: when every entry is older than the time window, the - count floor still keeps the last `limit * 2` entries.""" - cfg.conversation_history_limit = 3 # Keep last 3 pairs = 6 entries - cfg.conversation_history_window_minutes = 60 - agent = _make_agent(cfg, fake_backend) + @pytest.mark.asyncio + async def test_session_from_other_instructions_is_not_resumed( + self, cfg: MemclawConfig, fake_backend: FakeBackend, + ): + (cfg.memory_dir / "sessions.json").write_text(json.dumps({ + "cli": {"session_id": "old", "fingerprint": "stale", + "agents_hash": "x", "memory_hash": "y"}, + })) + agent = _chat_agent(cfg, fake_backend) + await agent.handle("hi") + + call = fake_backend.turn_calls[0] + assert call["resume_session_id"] is None + # A fresh session gets the files in its system prompt, not as updates. + assert "Updated AGENTS.md" not in call["context"] + stored = json.loads((cfg.memory_dir / "sessions.json").read_text()) + assert stored["cli"]["session_id"] == "sess-1" + agent.close() - # All 10 entries timestamped 2 hours ago (outside the window). - old_ts = (datetime.now() - timedelta(hours=2)).isoformat() - for i in range(10): - agent._history.append({ - "role": "user" if i % 2 == 0 else "assistant", - "content": f"message {i}", - "timestamp": old_ts, - }) + @pytest.mark.asyncio + async def test_backend_without_sessions_stores_nothing( + self, cfg: MemclawConfig, fake_backend: FakeBackend, + ): + fake_backend.session_id = None + agent = _chat_agent(cfg, fake_backend) + await agent.handle("hi") + assert not (cfg.memory_dir / "sessions.json").exists() + agent.close() + + @pytest.mark.asyncio + async def test_reset_conversation(self, cfg: MemclawConfig, fake_backend: FakeBackend): + agent = _chat_agent(cfg, fake_backend, platform="telegram") + await agent.handle("one", chat_id="42") + await agent.handle("other chat", chat_id="7") + agent.record_reminder_fired("⏰ Reminder: stretch", "42") + + await agent.reset_conversation("42") + + assert fake_backend.reset_calls == ["telegram:42"] + stored = json.loads((cfg.memory_dir / "sessions.json").read_text()) + assert "telegram:42" not in stored + assert "telegram:7" in stored + await agent.handle("two", chat_id="42") + call = fake_backend.turn_calls[-1] + assert call["resume_session_id"] is None + assert "stretch" not in call["context"] + agent.close() - agent._trim_history() - assert len(agent._history) == 6 - assert agent._history[0]["content"] == "message 4" - assert agent._history[-1]["content"] == "message 9" +class TestSystemPromptAndContext: + @pytest.mark.asyncio + async def test_system_prompt_is_static(self, cfg: MemclawConfig, fake_backend: FakeBackend): + cfg.agent_file.write_text("Be terse.") + cfg.memory_file.write_text("## Key Facts\n- Name: Test") + agent = _chat_agent(cfg, fake_backend) + + await agent.handle("one") + await agent.handle("two") + + first, second = (c["system_prompt"] for c in fake_backend.turn_calls) + assert first == second + assert "Be terse." in first + assert "- Name: Test" in first + assert "Current local time" not in first agent.close() - def test_history_window_dominates_floor( - self, cfg: MemclawConfig, fake_backend: FakeBackend + @pytest.mark.asyncio + async def test_context_block_carries_time_and_memories( + self, cfg: MemclawConfig, fake_backend: FakeBackend, ): - """When many recent entries fall inside the window, all of them are - kept even though that exceeds the count floor.""" - cfg.conversation_history_limit = 3 # floor = 6 entries - cfg.conversation_history_window_minutes = 60 - agent = _make_agent(cfg, fake_backend) + agent = _chat_agent(cfg, fake_backend) + agent.search.search = AsyncMock(return_value=[ + SearchResult(file_path="memory/2025-03-01.md", line_start=0, line_end=1, + content="Bought a red\nbike", score=0.9, match_type="hybrid"), + SearchResult(file_path=str(cfg.memory_file), line_start=0, line_end=1, + content="already in the system prompt", score=0.8, + match_type="hybrid"), + SearchResult(file_path="memory/2025-03-02.md", line_start=0, line_end=1, + content="x" * 2000, score=0.7, match_type="hybrid"), + ]) - recent_ts = (datetime.now() - timedelta(minutes=10)).isoformat() - for i in range(30): - agent._history.append({ - "role": "user" if i % 2 == 0 else "assistant", - "content": f"message {i}", - "timestamp": recent_ts, - }) + await agent.handle("what bike?") - agent._trim_history() + context = fake_backend.turn_calls[0]["context"] + assert context.startswith("\nCurrent local time: ") + assert context.endswith("") + assert "- [2025-03-01] Bought a red bike" in context + assert "already in the system prompt" not in context + assert "x" * 501 not in context # long snippets are cut + agent.close() - assert len(agent._history) == 30 - assert agent._history[0]["content"] == "message 0" - assert agent._history[-1]["content"] == "message 29" + @pytest.mark.asyncio + async def test_changed_agents_md_is_sent_once(self, cfg: MemclawConfig, fake_backend: FakeBackend): + cfg.agent_file.write_text("Be terse.") + agent = _chat_agent(cfg, fake_backend) + await agent.handle("one") + assert "Updated AGENTS.md" not in fake_backend.turn_calls[0]["context"] + + cfg.agent_file.write_text("Be terse.\nAlways answer in Spanish.") + await agent.handle("two") + await agent.handle("three") + + second = fake_backend.turn_calls[1]["context"] + assert "Updated AGENTS.md" in second + assert "Always answer in Spanish." in second + assert "Updated AGENTS.md" not in fake_backend.turn_calls[2]["context"] agent.close() - def test_history_union_keeps_floor_plus_window( - self, cfg: MemclawConfig, fake_backend: FakeBackend + @pytest.mark.asyncio + async def test_changed_memory_md_is_sent_once_even_after_restart( + self, cfg: MemclawConfig, fake_backend: FakeBackend, ): - """Mixed case: some entries within the window, some older. The retained - set is the union — at least the floor, plus any additional in-window - entries.""" - cfg.conversation_history_limit = 3 # floor = 6 entries - cfg.conversation_history_window_minutes = 60 - agent = _make_agent(cfg, fake_backend) + cfg.memory_file.write_text("## Key Facts\n- Old fact") + agent = _chat_agent(cfg, fake_backend) + await agent.handle("one") + agent.close() - old_ts = (datetime.now() - timedelta(hours=2)).isoformat() - recent_ts = (datetime.now() - timedelta(minutes=5)).isoformat() + cfg.memory_file.write_text("## Key Facts\n- New fact") + restarted = _chat_agent(cfg, fake_backend) + await restarted.handle("two") + await restarted.handle("three") - # 10 old entries, then 8 recent ones. - for i in range(10): - agent._history.append({ - "role": "user" if i % 2 == 0 else "assistant", - "content": f"old {i}", - "timestamp": old_ts, - }) - for i in range(8): - agent._history.append({ - "role": "user" if i % 2 == 0 else "assistant", - "content": f"recent {i}", - "timestamp": recent_ts, - }) + second = fake_backend.turn_calls[1]["context"] + assert "Updated permanent memory" in second + assert "- New fact" in second + assert "Updated permanent memory" not in fake_backend.turn_calls[2]["context"] + restarted.close() - agent._trim_history() + @pytest.mark.asyncio + async def test_own_writes_during_a_turn_are_not_sent_back( + self, cfg: MemclawConfig, fake_backend: FakeBackend, + ): + cfg.memory_file.write_text("## Key Facts\n- Old fact") + agent = _chat_agent(cfg, fake_backend) + + async def turn_that_saves(**kwargs): + # memory_save(permanent=true) / update_instructions mid-turn. + cfg.memory_file.write_text("## Key Facts\n- Old fact\n- Dog is Bruno") + cfg.agent_file.write_text("Be terse.") + return TurnResult(text="saved", session_id="sess-1") + + agent.backend.run_turn = turn_that_saves + await agent.handle("my dog is Bruno") + agent.backend.run_turn = FakeBackend.run_turn.__get__(fake_backend) + await agent.handle("next") + + context = fake_backend.turn_calls[-1]["context"] + assert "Updated permanent memory" not in context + assert "Updated AGENTS.md" not in context + agent.close() - # All 8 in-window entries are kept (exceeds floor of 6). - assert len(agent._history) == 8 - assert all(e["content"].startswith("recent") for e in agent._history) + @pytest.mark.asyncio + async def test_consolidation_during_a_turn_is_still_sent( + self, cfg: MemclawConfig, fake_backend: FakeBackend, + ): + cfg.memory_file.write_text("## Key Facts\n- Old fact") + agent = _chat_agent(cfg, fake_backend) + + async def turn_racing_consolidation(**kwargs): + cfg.memory_file.write_text("## Key Facts\n- Consolidated") + agent._memory_rewrites += 1 + return TurnResult(text="ok", session_id="sess-1") + + agent.backend.run_turn = turn_racing_consolidation + await agent.handle("one") + agent.backend.run_turn = FakeBackend.run_turn.__get__(fake_backend) + await agent.handle("two") + + context = fake_backend.turn_calls[-1]["context"] + assert "Updated permanent memory" in context + assert "- Consolidated" in context agent.close() - def test_history_floor_pulls_in_old_entries_to_meet_count( - self, cfg: MemclawConfig, fake_backend: FakeBackend + @pytest.mark.asyncio + async def test_reminder_notes_go_into_the_next_message_of_that_chat( + self, cfg: MemclawConfig, fake_backend: FakeBackend, ): - """When the time window holds fewer than the floor count, the floor - pulls in older entries to bring the total up.""" - cfg.conversation_history_limit = 3 # floor = 6 entries - cfg.conversation_history_window_minutes = 60 - agent = _make_agent(cfg, fake_backend) + agent = _chat_agent(cfg, fake_backend, platform="telegram") + agent.record_reminder_fired("⏰ Reminder: call Alex", "42") - old_ts = (datetime.now() - timedelta(hours=2)).isoformat() - recent_ts = (datetime.now() - timedelta(minutes=5)).isoformat() - - # 50 old entries + 4 recent — only 4 are in-window but floor is 6. - for i in range(50): - agent._history.append({ - "role": "user" if i % 2 == 0 else "assistant", - "content": f"old {i}", - "timestamp": old_ts, - }) - for i in range(4): - agent._history.append({ - "role": "user" if i % 2 == 0 else "assistant", - "content": f"recent {i}", - "timestamp": recent_ts, - }) - - agent._trim_history() - - # Keep exactly the floor (6): 2 trailing old + all 4 recent. - assert len(agent._history) == 6 - assert agent._history[0]["content"] == "old 48" - assert agent._history[1]["content"] == "old 49" - assert agent._history[-1]["content"] == "recent 3" - agent.close() - - def test_image_placeholder_in_history(self, cfg: MemclawConfig, fake_backend: FakeBackend): - """When an image is sent, history should store a placeholder, not base64.""" - agent = _make_agent(cfg, fake_backend) + await agent.handle("other chat", chat_id="7") + await agent.handle("done!", chat_id="42") + await agent.handle("and now?", chat_id="42") + + other, first, second = (c["context"] for c in fake_backend.turn_calls) + assert "call Alex" not in other + assert "- ⏰ Reminder: call Alex" in first + assert "call Alex" not in second + agent.close() + + @pytest.mark.asyncio + async def test_reminder_note_survives_a_failed_turn( + self, cfg: MemclawConfig, fake_backend: FakeBackend, + ): + agent = _chat_agent(cfg, fake_backend, platform="telegram") + agent.record_reminder_fired("⏰ Reminder: call Alex", "42") + fake_backend.turn_error = RuntimeError("boom") + with pytest.raises(RuntimeError): + await agent.handle("hi", chat_id="42") + fake_backend.turn_error = None + await agent.handle("hi again", chat_id="42") + assert "call Alex" in fake_backend.turn_calls[-1]["context"] + agent.close() - # Simulate what handle() does for images - image_b64 = "base64data..." - history_content = "[User sent a photo]" if image_b64 else "text" - agent._history.append({ - "role": "user", - "content": history_content, - "timestamp": datetime.now().isoformat(), - }) - assert agent._history[0]["content"] == "[User sent a photo]" - assert "base64" not in agent._history[0]["content"] +class TestBackgroundConsolidation: + @pytest.mark.asyncio + async def test_runs_after_the_reply_without_blocking_it( + self, cfg: MemclawConfig, fake_backend: FakeBackend, + ): + agent = _chat_agent(cfg, fake_backend) + started = asyncio.Event() + release = asyncio.Event() + + async def slow_consolidate(**kwargs): + started.set() + await release.wait() + return True + + agent._maybe_consolidate = AsyncMock(side_effect=slow_consolidate) + + await agent.handle("hi") # returns while consolidation is still pending + await asyncio.wait_for(started.wait(), 1) + await agent.handle("again") # one already running: no second one + release.set() + await agent._consolidation_task + assert agent._maybe_consolidate.await_count == 1 + agent.close() + + @pytest.mark.asyncio + async def test_failure_is_logged_not_raised(self, cfg: MemclawConfig, fake_backend: FakeBackend): + agent = _chat_agent(cfg, fake_backend) + agent._maybe_consolidate = AsyncMock(side_effect=RuntimeError("API down")) + reply, _ = await agent.handle("hi") + await agent._consolidation_task # does not raise + assert reply agent.close() + @pytest.mark.asyncio + async def test_aclose_cancels_a_running_consolidation( + self, cfg: MemclawConfig, fake_backend: FakeBackend, + ): + agent = _chat_agent(cfg, fake_backend) + async def forever(**kwargs): + await asyncio.sleep(60) + + agent._maybe_consolidate = AsyncMock(side_effect=forever) + await agent.handle("hi") + await asyncio.sleep(0) + task = agent._consolidation_task + await agent.aclose() + assert task.cancelled() + # ──────────────────────────────────────────────────────────────────── # Spec #2: Memory Consolidation @@ -369,55 +534,29 @@ async def test_content_limit_30000_chars(self, cfg: MemclawConfig, fake_backend: # ──────────────────────────────────────────────────────────────────── -# Spec #3: MEMORY.md Context Strategy +# Spec #3: MEMORY.md in the system prompt # ──────────────────────────────────────────────────────────────────── -class TestContextStrategy: - @pytest.mark.asyncio - async def test_small_memory_included_in_full(self, cfg: MemclawConfig, fake_backend: FakeBackend): - """MEMORY.md under 4000 chars should be included completely.""" - agent = _make_agent(cfg, fake_backend) - - small_content = "## Key Facts\n\n- I like Python\n- My name is Test" - cfg.memory_file.write_text(small_content) - - agent.search.search = AsyncMock(return_value=[]) +class TestPermanentMemoryPrompt: + def test_small_memory_included_in_full(self, cfg: MemclawConfig): + from memclaw.agent import _load_permanent_memory - context = await agent.build_context("hello") - assert small_content in context - agent.close() + cfg.memory_file.write_text("## Key Facts\n\n- I like Python") + assert _load_permanent_memory(cfg) == "## Key Facts\n\n- I like Python" - @pytest.mark.asyncio - async def test_large_memory_truncated_with_search(self, cfg: MemclawConfig, fake_backend: FakeBackend): - """MEMORY.md over 4000 chars: first 2000 + semantic search results.""" - agent = _make_agent(cfg, fake_backend) + def test_runaway_memory_is_truncated(self, cfg: MemclawConfig): + from memclaw.agent import _MEMORY_PROMPT_CHARS, _load_permanent_memory - large_content = "## Key Facts\n\n" + "Important fact. " * 400 - cfg.memory_file.write_text(large_content) + cfg.memory_file.write_text("x" * (_MEMORY_PROMPT_CHARS + 100)) + text = _load_permanent_memory(cfg) + assert len(text) < _MEMORY_PROMPT_CHARS + 100 + assert "truncated" in text - memory_chunk = SearchResult( - file_path=str(cfg.memory_file), - line_start=100, - line_end=110, - content="Relevant chunk from MEMORY.md about Python", - score=0.8, - match_type="vector", - ) - agent.search.search = AsyncMock(side_effect=[ - [memory_chunk], # file_filter="MEMORY.md" call - [], # general search call - ]) + def test_empty_memory(self, cfg: MemclawConfig): + from memclaw.agent import _load_permanent_memory - context = await agent.build_context("tell me about Python") - - assert large_content[:100] in context - assert "Relevant chunk from MEMORY.md about Python" in context - assert len(context) < len(large_content) - - calls = agent.search.search.call_args_list - assert calls[0].kwargs.get("file_filter") == "MEMORY.md" - - agent.close() + cfg.memory_file.write_text("") + assert _load_permanent_memory(cfg) == "(empty)" # ──────────────────────────────────────────────────────────────────── diff --git a/tests/test_config.py b/tests/test_config.py index d5943b6..be3097e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,14 +1,9 @@ -"""Tests for new MemclawConfig fields (specs #1, #2, #4, #5).""" +"""Tests for MemclawConfig fields (specs #2, #4, #5).""" from pathlib import Path from memclaw.config import MemclawConfig -def test_default_conversation_history_limit(tmp_path: Path): - cfg = MemclawConfig(memory_dir=tmp_path / "m", openai_api_key="k", anthropic_api_key="k") - assert cfg.conversation_history_limit == 10 - - def test_default_consolidation_threshold(tmp_path: Path): cfg = MemclawConfig(memory_dir=tmp_path / "m", openai_api_key="k", anthropic_api_key="k") assert cfg.consolidation_threshold == 7 @@ -29,12 +24,10 @@ def test_custom_values(tmp_path: Path): memory_dir=tmp_path / "m", openai_api_key="k", anthropic_api_key="k", - conversation_history_limit=5, consolidation_threshold=3, decay_half_life_days=60, mmr_lambda=0.5, ) - assert cfg.conversation_history_limit == 5 assert cfg.consolidation_threshold == 3 assert cfg.decay_half_life_days == 60 assert cfg.mmr_lambda == 0.5 @@ -77,3 +70,22 @@ def test_claude_cli_model_vars_are_not_picked_up( monkeypatch.setenv("ANTHROPIC_MODEL", "claude-opus-5") cfg = MemclawConfig(memory_dir=tmp_path / "m", openai_api_key="k", anthropic_api_key="k") assert cfg.claude_model == "" + + +def test_rolling_history_settings_are_gone(tmp_path: Path): + """Conversations live in the backend session now; there is no window.""" + cfg = MemclawConfig(memory_dir=tmp_path / "m", openai_api_key="k", anthropic_api_key="k") + assert not hasattr(cfg, "conversation_history_limit") + assert not hasattr(cfg, "conversation_history_window_minutes") + + +def test_update_env_file_merges_and_drops(tmp_path: Path, monkeypatch): + from memclaw import setup + + env_file = tmp_path / ".env" + env_file.write_text("OPENAI_API_KEY=k\nCLAUDE_MODEL=old\nCLAUDE_EFFORT=high\n") + monkeypatch.setattr(setup, "ENV_FILE", env_file) + + setup.update_env_file({"CLAUDE_MODEL": "claude-haiku-4-5"}, drop_keys=["CLAUDE_EFFORT"]) + + assert env_file.read_text() == "OPENAI_API_KEY=k\nCLAUDE_MODEL=claude-haiku-4-5\n" diff --git a/tests/test_handlers.py b/tests/test_handlers.py index f5277bd..41a5630 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -210,3 +210,203 @@ def ev(*, is_group: bool, is_from_me: bool, chat_user: str, sender_user: str): assert check(None, out_to_friend) is False, "outgoing DM to friend must be ignored" assert check(None, in_from_friend) is False assert check(None, group) is False + + +# ──────────────────────────────────────────────────────────────────── +# /new, /model, /effort +# ──────────────────────────────────────────────────────────────────── + +def _telegram_handlers(*, allowed: bool = True): + from memclaw.bot.handlers import MessageHandlers + + with patch.object(MessageHandlers, "__init__", lambda self, *a, **kw: None): + handlers = MessageHandlers.__new__(MessageHandlers) + handlers.agent = MagicMock() + handlers.agent.reset_conversation = AsyncMock() + handlers.config = MagicMock() + handlers.config.allowed_user_ids_list = [1] if allowed else [] + return handlers + + +def _command_update(chat_id: int = 99): + update = MagicMock() + update.effective_user.id = 1 + update.effective_chat.id = chat_id + update.message.reply_text = AsyncMock() + return update + + +class _SelectableBackend: + """Minimal backend exposing the model-selection surface.""" + + display_name = "Fake" + supports_model_selection = True + + def __init__(self): + from memclaw.backends.claude_models import ModelInfo + + self.models = [ + ModelInfo("claude-opus-5", "Claude Opus 5", "", ["low", "medium", "high"]), + ModelInfo("claude-haiku-4-5", "Claude Haiku 4.5", "", []), + ] + self.model = "claude-opus-5" + self.effort = "high" + + async def list_models(self): + return self.models + + async def effort_levels(self): + info = next((m for m in self.models if m.id == self.model), None) + return info.effort_levels if info else None + + async def set_model(self, model_id): + if model_id not in [m.id for m in self.models]: + raise ValueError(f"Unknown model '{model_id}'.") + self.model = model_id + if model_id == "claude-haiku-4-5": + self.effort = None + return "Claude Haiku 4.5 doesn't support an effort setting, so none is used." + return "" + + async def set_effort(self, level): + if level not in await self.effort_levels(): + raise ValueError(f"The current model doesn't support '{level}'.") + self.effort = level + + +class TestNewCommand: + @pytest.mark.asyncio + async def test_resets_the_chat(self): + handlers = _telegram_handlers() + update = _command_update(chat_id=99) + await handlers.new_command(update, MagicMock()) + handlers.agent.reset_conversation.assert_awaited_once_with("99") + update.message.reply_text.assert_awaited_once() + + @pytest.mark.asyncio + async def test_ignores_unknown_users(self): + handlers = _telegram_handlers(allowed=False) + update = _command_update() + await handlers.new_command(update, MagicMock()) + handlers.agent.reset_conversation.assert_not_awaited() + update.message.reply_text.assert_not_awaited() + + @pytest.mark.asyncio + async def test_slack_new_message_resets_the_channel(self): + from memclaw.bot.slack_handlers import SlackHandlers + + with patch.object(SlackHandlers, "__init__", lambda self, *a, **kw: None): + handlers = SlackHandlers.__new__(SlackHandlers) + handlers.config = MagicMock() + handlers.config.slack_allowed_channels_list = [] + handlers.config.slack_allowed_users_list = [] + handlers.agent = MagicMock() + handlers.agent.reset_conversation = AsyncMock() + handlers.agent.handle = AsyncMock() + say = AsyncMock() + + await handlers._route_event( + {"channel": "C1", "user": "U1", "text": "<@UBOT> /new", "ts": "1.0"}, + say, MagicMock(), + ) + + handlers.agent.reset_conversation.assert_awaited_once_with("C1") + handlers.agent.handle.assert_not_awaited() + say.assert_awaited_once() + + +class TestModelCommand: + @pytest.mark.asyncio + async def test_lists_models_numbered(self): + from memclaw.bot.handlers import model_command_reply + + reply = await model_command_reply(_SelectableBackend(), []) + assert "Model: claude-opus-5" in reply + assert "Effort: high" in reply + assert "1. Claude Opus 5 - claude-opus-5 (current)" in reply + assert "2. Claude Haiku 4.5 - claude-haiku-4-5" in reply + + @pytest.mark.asyncio + async def test_switch_by_number_reports_effort_change(self): + from memclaw.bot.handlers import model_command_reply + + backend = _SelectableBackend() + reply = await model_command_reply(backend, ["2"]) + assert backend.model == "claude-haiku-4-5" + assert "Switched to claude-haiku-4-5." in reply + assert "doesn't support an effort" in reply + assert "conversation is kept" in reply + + @pytest.mark.asyncio + async def test_switch_by_id(self): + from memclaw.bot.handlers import model_command_reply + + backend = _SelectableBackend() + backend.model = "claude-haiku-4-5" + await model_command_reply(backend, ["claude-opus-5"]) + assert backend.model == "claude-opus-5" + + @pytest.mark.asyncio + async def test_bad_number_and_unknown_id(self): + from memclaw.bot.handlers import model_command_reply + + backend = _SelectableBackend() + assert "no model number 9" in await model_command_reply(backend, ["9"]) + assert "Unknown model" in await model_command_reply(backend, ["claude-nope"]) + assert backend.model == "claude-opus-5" + + @pytest.mark.asyncio + async def test_unsupported_backend(self): + from memclaw.bot.handlers import effort_command_reply, model_command_reply + + backend = MagicMock(supports_model_selection=False, display_name="Cursor SDK") + assert "isn't supported for the Cursor SDK backend" in await model_command_reply(backend, []) + assert "isn't supported for the Cursor SDK backend" in await effort_command_reply(backend, ["low"]) + + @pytest.mark.asyncio + async def test_telegram_handler_passes_args(self): + handlers = _telegram_handlers() + handlers.agent.backend = _SelectableBackend() + update = _command_update() + context = MagicMock() + context.args = ["2"] + await handlers.model_command(update, context) + assert handlers.agent.backend.model == "claude-haiku-4-5" + update.message.reply_text.assert_awaited_once() + + +class TestEffortCommand: + @pytest.mark.asyncio + async def test_shows_supported_levels(self): + from memclaw.bot.handlers import effort_command_reply + + reply = await effort_command_reply(_SelectableBackend(), []) + assert "Effort: high" in reply + assert "low, medium, high" in reply + + @pytest.mark.asyncio + async def test_model_without_effort(self): + from memclaw.bot.handlers import effort_command_reply + + backend = _SelectableBackend() + backend.model, backend.effort = "claude-haiku-4-5", None + reply = await effort_command_reply(backend, []) + assert "doesn't support an effort setting" in reply + + @pytest.mark.asyncio + async def test_sets_level(self): + from memclaw.bot.handlers import effort_command_reply + + backend = _SelectableBackend() + reply = await effort_command_reply(backend, ["low"]) + assert backend.effort == "low" + assert "Effort set to low." in reply + + @pytest.mark.asyncio + async def test_rejected_level(self): + from memclaw.bot.handlers import effort_command_reply + + backend = _SelectableBackend() + reply = await effort_command_reply(backend, ["max"]) + assert "doesn't support 'max'" in reply + assert backend.effort == "high" diff --git a/uv.lock b/uv.lock index f62475e..c999ef0 100644 --- a/uv.lock +++ b/uv.lock @@ -410,7 +410,7 @@ wheels = [ [[package]] name = "claude-agent-sdk" -version = "0.2.152" +version = "0.2.164" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -419,13 +419,13 @@ dependencies = [ { name = "sniffio" }, { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/49/61/de2194d266fcede2a737996e800e7726079730f81ef03d41117e7888be04/claude_agent_sdk-0.2.152.tar.gz", hash = "sha256:0822bbdf700ccdd1acecf8365345a56bde07833afe7f8bd1bee68cba0e2fcf9a", size = 344710, upload-time = "2026-09-02T22:47:58.633Z" } +sdist = { url = "https://files.pythonhosted.org/packages/72/f2/7a954ff5180f50cd1786a05e9e0d3a53ddfe416b48a05ac6c6f20d0bd0ab/claude_agent_sdk-0.2.164.tar.gz", hash = "sha256:eeb967a43c3d1e5589b95ca3c050ec45ce72e242bfab1ded2205b33c61c8b430", size = 365170, upload-time = "2026-10-06T23:50:21.704Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a0/d3/ee2b674958e105cf26a2eb4d512c46b02fcfdc8fadbb14dde750bbfaf9b1/claude_agent_sdk-0.2.152-py3-none-macosx_11_0_arm64.whl", hash = "sha256:6cc1e949743c7634274ae526bfd1c0a216ff9d71c6b36d94cab80f32b0d32b4c", size = 85989333, upload-time = "2026-09-02T22:48:02.484Z" }, - { url = "https://files.pythonhosted.org/packages/92/c4/6564671e6d303012b5c0bd6650b9977317c3b5c83914cf5ac62068fa85e6/claude_agent_sdk-0.2.152-py3-none-macosx_11_0_x86_64.whl", hash = "sha256:9c79c4ef91b15e5fae6aba62a82afb983dc360bc159902c541046227e2852c86", size = 90503485, upload-time = "2026-09-02T22:48:06.357Z" }, - { url = "https://files.pythonhosted.org/packages/13/5d/1efd9e3332d2909e015c6571a4b5ed937e9126b9c702217813247c1c8052/claude_agent_sdk-0.2.152-py3-none-manylinux_2_17_aarch64.whl", hash = "sha256:1e5b619e6fa3c98f06ddc107652f7820c4573fb73a9a24263568d9974078bbc2", size = 95850705, upload-time = "2026-09-02T22:48:10.493Z" }, - { url = "https://files.pythonhosted.org/packages/95/a0/f921e52fdbfd400e46fbc1e707e6004a3ef99ace6cb4555429a3a8d3f371/claude_agent_sdk-0.2.152-py3-none-manylinux_2_17_x86_64.whl", hash = "sha256:f885bd935054082a599dcf92bbdaf3c0b853117a79375cb3aae74674a8124326", size = 96039825, upload-time = "2026-09-02T22:48:14.587Z" }, - { url = "https://files.pythonhosted.org/packages/e6/be/dd113754761c6eef6eec1f2b148c825f925d4bfaa3f946656925239c54f9/claude_agent_sdk-0.2.152-py3-none-win_amd64.whl", hash = "sha256:f422c256358997ffa3a740f9fdc8127aa8404a9d57493502dc7d2ac05ecb5486", size = 98656506, upload-time = "2026-09-02T22:48:19.109Z" }, + { url = "https://files.pythonhosted.org/packages/80/eb/6e986387fbab854040c9c9498dcce6ff6e6d882c03059caae53f45460dbd/claude_agent_sdk-0.2.164-py3-none-macosx_11_0_arm64.whl", hash = "sha256:b810e357691aa54ae34739fdeac747819707099903e0621f31989354b9153f05", size = 97424192, upload-time = "2026-10-06T23:50:25.835Z" }, + { url = "https://files.pythonhosted.org/packages/52/7f/239855350c0979fd522d5dd2d85e7fb3f6cec6bb8b50c6ef9e6aca2fbc9f/claude_agent_sdk-0.2.164-py3-none-macosx_11_0_x86_64.whl", hash = "sha256:3c9523749c4163c990b937b7408907f8a4c5dc09123ad9eba7150630b5c2c86a", size = 101920894, upload-time = "2026-10-06T23:50:29.728Z" }, + { url = "https://files.pythonhosted.org/packages/fe/81/d06637d9344a3552619cf5ba990d09f068a321d0adc24b6b79e71af38508/claude_agent_sdk-0.2.164-py3-none-manylinux_2_17_aarch64.whl", hash = "sha256:1a24aa955dac8a231c1d782d89cb2cae8d875159398317fbaa8cd777cb94c311", size = 106886299, upload-time = "2026-10-06T23:50:36.857Z" }, + { url = "https://files.pythonhosted.org/packages/bf/07/a9ff8cd09729e764802013162306fd7356d6134d04421fc424795aae3290/claude_agent_sdk-0.2.164-py3-none-manylinux_2_17_x86_64.whl", hash = "sha256:6e45da8ffe8bf3b0472e291584af4ee40d315e02c3f1aa7b4cc86b5a0425f552", size = 107239153, upload-time = "2026-10-06T23:50:40.974Z" }, + { url = "https://files.pythonhosted.org/packages/32/2d/3ab6918baeca42456ec2d431a05656dbb1b4f855c5b1eac2aa60cdd98473/claude_agent_sdk-0.2.164-py3-none-win_amd64.whl", hash = "sha256:d425982a50f804a28b538248fed96a74778beaca5c14830d0b5a52d1c64d511f", size = 109850011, upload-time = "2026-10-06T23:50:45.287Z" }, ] [[package]] @@ -1004,7 +1004,7 @@ dev = [ requires-dist = [ { name = "aiohttp", specifier = ">=3.9.0" }, { name = "beautifulsoup4", specifier = ">=4.12.0" }, - { name = "claude-agent-sdk", specifier = ">=0.2.139,<0.3" }, + { name = "claude-agent-sdk", specifier = ">=0.2.163,<0.3" }, { name = "click", specifier = ">=8.0.0" }, { name = "curl-cffi", specifier = ">=0.7.0" }, { name = "cursor-sdk", marker = "extra == 'cursor'", specifier = ">=0.1.0" }, From a7114cef663528ff72da1ef9a1215d2d39d2921f Mon Sep 17 00:00:00 2001 From: Emanuilo Jovanovic Date: Wed, 7 Oct 2026 21:36:56 +0200 Subject: [PATCH 02/10] feat: inline keyboard menus for /model and /effort --- README.md | 2 +- memclaw/bot/handlers.py | 223 ++++++++++++++++++++++++++++------------ memclaw/cli.py | 14 ++- tests/test_handlers.py | 182 ++++++++++++++++++++++---------- 4 files changed, 296 insertions(+), 125 deletions(-) diff --git a/README.md b/README.md index 0b429d1..f0a04aa 100644 --- a/README.md +++ b/README.md @@ -76,7 +76,7 @@ Each chat is one ongoing conversation that survives restarts. Send `/new` to sta The Telegram bot shows a **typing indicator** while processing so you know it's working on your request. -Commands: `/new` starts a fresh conversation, `/model` lists the available models (`/model ` switches), and `/effort` shows or sets the reasoning effort. Model and effort changes are saved to `~/.memclaw/.env` and keep the current conversation. They're available with the Claude backend. +Commands: `/new` starts a fresh conversation, `/model` shows the available models as buttons (tap one to switch, or send `/model `), and `/effort` does the same for the reasoning effort. Model and effort changes are saved to `~/.memclaw/.env` and keep the current conversation. They're available with the Claude backend. #### Setup diff --git a/memclaw/bot/handlers.py b/memclaw/bot/handlers.py index c043244..ed4e8e9 100644 --- a/memclaw/bot/handlers.py +++ b/memclaw/bot/handlers.py @@ -11,8 +11,9 @@ from loguru import logger from openai import AsyncOpenAI -from telegram import Update +from telegram import InlineKeyboardButton, InlineKeyboardMarkup, Update from telegram.constants import ChatAction +from telegram.error import BadRequest from telegram.ext import ContextTypes from ..agent import MemclawAgent @@ -43,80 +44,116 @@ async def _typing_loop(bot, chat_id: int): ("effort", "Show or change the reasoning effort"), ] +# Telegram limits callback_data to 64 bytes. +CALLBACK_DATA_MAX_BYTES = 64 -async def model_command_reply(backend: AgentBackend, args: list[str]) -> str: - """Reply for `/model` (list models) or `/model ` (switch).""" - if not backend.supports_model_selection: - return f"Switching models isn't supported for the {backend.display_name} backend yet." +_APPLIES = "Applies from your next message; the conversation is kept." + + +def model_callback_data(index: int, model_id: str) -> str: + """``m:``, or ``mi:`` if the id would exceed Telegram's limit.""" + data = f"m:{model_id}" + return data if len(data.encode()) <= CALLBACK_DATA_MAX_BYTES else f"mi:{index}" + + +def _unsupported(backend: AgentBackend, what: str) -> str: + return f"{what} isn't supported for the {backend.display_name} backend yet." + + +async def _load_models(backend: AgentBackend) -> list: try: - models = await backend.list_models() + return await backend.list_models() except Exception as exc: logger.warning("Could not fetch the model list: {exc}", exc=exc) - models = [] - - arg = " ".join(args).strip() - if arg: - model_id = arg - if arg.isdigit(): - n = int(arg) - if not 1 <= n <= len(models): - return f"There's no model number {n}. Send /model to see the list." - model_id = models[n - 1].id - try: - note = await backend.set_model(model_id) - except ValueError as exc: - return f"{exc} Send /model to see the available models." - lines = [f"Switched to {backend.model}."] - if note: - lines.append(note) - lines.append(f"Effort: {backend.effort or 'model default'}") - lines.append("Applies from your next message; the conversation is kept.") - return "\n".join(lines) - - lines = [ - f"Model: {backend.model}", - f"Effort: {backend.effort or 'model default'}", - "", - ] - if models: - lines.append("Available models:") - for i, m in enumerate(models, 1): - current = " (current)" if m.id == backend.model else "" - lines.append(f"{i}. {m.display_name} - {m.id}{current}") - lines += ["", "Send /model to switch."] + return [] + + +def _status_lines(backend: AgentBackend, models: list, levels: list[str] | None) -> list[str]: + info = next((m for m in models if m.id == backend.model), None) + model = f"{info.display_name} ({info.id})" if info else backend.model + if levels == []: + effort = "not supported by this model" else: - lines.append("Couldn't load the model list right now; send /model to switch.") - return "\n".join(lines) + effort = backend.effort or "model default" + return [f"Model: {model}", f"Effort: {effort}"] -async def effort_command_reply(backend: AgentBackend, args: list[str]) -> str: - """Reply for `/effort` (show levels) or `/effort ` (set it).""" +async def model_menu( + backend: AgentBackend, note: str = "", +) -> tuple[str, InlineKeyboardMarkup | None]: + """`/model` reply: current settings plus one button per model.""" if not backend.supports_model_selection: - return f"Changing the effort isn't supported for the {backend.display_name} backend yet." - arg = " ".join(args).strip() - if arg: - try: - await backend.set_effort(arg) - except ValueError as exc: - return str(exc) - return ( - f"Effort set to {backend.effort}.\n" - "Applies from your next message; the conversation is kept." - ) - + return _unsupported(backend, "Switching models"), None + models = await _load_models(backend) + lines = _status_lines(backend, models, await backend.effort_levels()) + if note: + lines += ["", note] + lines.append("") + if not models: + lines.append("Couldn't load the model list right now; send /model to switch.") + return "\n".join(lines), None + lines.append("Tap a model to switch, or send /model .") + keyboard = InlineKeyboardMarkup([ + [InlineKeyboardButton( + ("✓ " if m.id == backend.model else "") + m.display_name, + callback_data=model_callback_data(i, m.id), + )] + for i, m in enumerate(models) + ]) + return "\n".join(lines), keyboard + + +async def effort_menu( + backend: AgentBackend, note: str = "", +) -> tuple[str, InlineKeyboardMarkup | None]: + """`/effort` reply: current settings plus a row of level buttons.""" + if not backend.supports_model_selection: + return _unsupported(backend, "Changing the effort"), None levels = await backend.effort_levels() - lines = [ - f"Model: {backend.model}", - f"Effort: {backend.effort or 'model default'}", - "", - ] + lines = _status_lines(backend, await _load_models(backend), levels) + if note: + lines += ["", note] + lines.append("") if levels == []: lines.append("This model doesn't support an effort setting.") - else: - if levels: - lines.append(f"Levels this model supports: {', '.join(levels)}") + return "\n".join(lines), None + if levels is None: lines.append("Send /effort to change it.") - return "\n".join(lines) + return "\n".join(lines), None + lines.append("Tap a level, or send /effort .") + keyboard = InlineKeyboardMarkup([[ + InlineKeyboardButton( + ("✓ " if lv == backend.effort else "") + lv, callback_data=f"e:{lv}", + ) + for lv in levels + ]]) + return "\n".join(lines), keyboard + + +async def select_model(backend: AgentBackend, model_id: str) -> str: + """Switch to *model_id*. Returns a note for the user; raises ValueError + for an unknown model.""" + if model_id == backend.model: + return f"Already using {model_id}." + note = await backend.set_model(model_id) + return " ".join(p for p in (f"Switched to {backend.model}.", note, _APPLIES) if p) + + +async def select_effort(backend: AgentBackend, level: str) -> str: + """Set the effort level. Returns a note; raises ValueError if rejected.""" + await backend.set_effort(level) + return f"Effort set to {backend.effort}. {_APPLIES}" + + +async def _resolve_model_arg(backend: AgentBackend, arg: str) -> str: + """`/model ` → model id. Raises ValueError.""" + if not arg.isdigit(): + return arg + models = await _load_models(backend) + n = int(arg) + if not 1 <= n <= len(models): + raise ValueError(f"There's no model number {n}.") + return models[n - 1].id class MessageHandlers: @@ -221,14 +258,66 @@ async def new_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE): async def model_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE): if not self._check_user(update.effective_user.id): return - reply = await model_command_reply(self.agent.backend, list(context.args or [])) - await update.message.reply_text(reply) + backend = self.agent.backend + arg = " ".join(context.args or []).strip() + note = "" + if arg and backend.supports_model_selection: + try: + note = await select_model(backend, await _resolve_model_arg(backend, arg)) + except ValueError as exc: + note = f"{exc} Pick one below." + text, keyboard = await model_menu(backend, note) + await update.message.reply_text(text, reply_markup=keyboard) async def effort_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE): if not self._check_user(update.effective_user.id): return - reply = await effort_command_reply(self.agent.backend, list(context.args or [])) - await update.message.reply_text(reply) + backend = self.agent.backend + arg = " ".join(context.args or []).strip() + note = "" + if arg and backend.supports_model_selection: + try: + note = await select_effort(backend, arg) + except ValueError as exc: + note = str(exc) + text, keyboard = await effort_menu(backend, note) + await update.message.reply_text(text, reply_markup=keyboard) + + async def handle_callback(self, update: Update, context: ContextTypes.DEFAULT_TYPE): + """Button presses on the /model and /effort menus.""" + query = update.callback_query + if query is None: + return + if not self._check_user(update.effective_user.id): + await query.answer("You're not allowed to use this bot.") + return + backend = self.agent.backend + kind, _, value = (query.data or "").partition(":") + if kind not in ("m", "mi", "e") or not backend.supports_model_selection: + await query.answer() + return + try: + if kind == "e": + note = await select_effort(backend, value) + await query.answer(f"Effort: {backend.effort}") + text, keyboard = await effort_menu(backend, note) + else: + model_id = value + if kind == "mi": + models = await _load_models(backend) + if not (value.isdigit() and int(value) < len(models)): + raise ValueError("That model is no longer available; send /model again.") + model_id = models[int(value)].id + note = await select_model(backend, model_id) + await query.answer(f"Model: {backend.model}") + text, keyboard = await model_menu(backend, note) + except ValueError as exc: + await query.answer(str(exc)[:200], show_alert=True) + return + try: + await query.edit_message_text(text, reply_markup=keyboard) + except BadRequest as exc: # e.g. "message is not modified" + logger.debug("Editing the menu after a button press failed: {exc}", exc=exc) # ------------------------------------------------------------------ # Message handlers — everything goes through the agent diff --git a/memclaw/cli.py b/memclaw/cli.py index 1c630b3..844a7e9 100644 --- a/memclaw/cli.py +++ b/memclaw/cli.py @@ -185,7 +185,13 @@ def _run_telegram(config: MemclawConfig) -> None: from openai import AsyncOpenAI from telegram import BotCommand from telegram.error import NetworkError, TelegramError, TimedOut - from telegram.ext import Application, CommandHandler, MessageHandler, filters + from telegram.ext import ( + Application, + CallbackQueryHandler, + CommandHandler, + MessageHandler, + filters, + ) from .bot.handlers import BOT_COMMANDS, MessageHandlers @@ -242,6 +248,9 @@ async def _model(update, context): async def _effort(update, context): await context.bot_data["handlers"].effort_command(update, context) + async def _callback(update, context): + await context.bot_data["handlers"].handle_callback(update, context) + async def _text(update, context): await context.bot_data["handlers"].handle_text(update, context) @@ -263,6 +272,7 @@ async def _on_error(update, context): app.add_handler(CommandHandler("new", _new)) app.add_handler(CommandHandler("model", _model)) app.add_handler(CommandHandler("effort", _effort)) + app.add_handler(CallbackQueryHandler(_callback)) app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, _text)) app.add_handler(MessageHandler(filters.PHOTO, _photo)) app.add_handler(MessageHandler(filters.VOICE, _voice)) @@ -273,7 +283,7 @@ async def _on_error(update, context): f"[green]Starting Memclaw Telegram bot...[/green] " f"(allowed users: {masked})" ) - app.run_polling(allowed_updates=["message"]) + app.run_polling(allowed_updates=["message", "callback_query"]) def _run_whatsapp(config: MemclawConfig) -> None: diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 41a5630..27dfaaf 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -315,98 +315,170 @@ async def test_slack_new_message_resets_the_channel(self): say.assert_awaited_once() +def _buttons(keyboard): + return [(b.text, b.callback_data) for row in keyboard.inline_keyboard for b in row] + + +def _callback_update(data: str, *, allowed_user: int = 1): + update = MagicMock() + update.effective_user.id = allowed_user + update.callback_query.data = data + update.callback_query.answer = AsyncMock() + update.callback_query.edit_message_text = AsyncMock() + return update + + class TestModelCommand: @pytest.mark.asyncio - async def test_lists_models_numbered(self): - from memclaw.bot.handlers import model_command_reply + async def test_menu_has_a_button_per_model(self): + from memclaw.bot.handlers import model_menu + + text, keyboard = await model_menu(_SelectableBackend()) + assert "Model: Claude Opus 5 (claude-opus-5)" in text + assert "Effort: high" in text + assert _buttons(keyboard) == [ + ("✓ Claude Opus 5", "m:claude-opus-5"), + ("Claude Haiku 4.5", "m:claude-haiku-4-5"), + ] - reply = await model_command_reply(_SelectableBackend(), []) - assert "Model: claude-opus-5" in reply - assert "Effort: high" in reply - assert "1. Claude Opus 5 - claude-opus-5 (current)" in reply - assert "2. Claude Haiku 4.5 - claude-haiku-4-5" in reply + def test_long_model_id_falls_back_to_index(self): + from memclaw.bot.handlers import model_callback_data + + assert model_callback_data(3, "x" * 70) == "mi:3" + assert model_callback_data(3, "claude-opus-5") == "m:claude-opus-5" @pytest.mark.asyncio - async def test_switch_by_number_reports_effort_change(self): - from memclaw.bot.handlers import model_command_reply + async def test_button_press_switches_and_edits_the_menu(self): + handlers = _telegram_handlers() + handlers.agent.backend = backend = _SelectableBackend() + update = _callback_update("m:claude-haiku-4-5") + + await handlers.handle_callback(update, MagicMock()) - backend = _SelectableBackend() - reply = await model_command_reply(backend, ["2"]) assert backend.model == "claude-haiku-4-5" - assert "Switched to claude-haiku-4-5." in reply - assert "doesn't support an effort" in reply - assert "conversation is kept" in reply + update.callback_query.answer.assert_awaited_once() + text = update.callback_query.edit_message_text.await_args.args[0] + assert "Switched to claude-haiku-4-5." in text + assert "doesn't support an effort" in text + assert "conversation is kept" in text + keyboard = update.callback_query.edit_message_text.await_args.kwargs["reply_markup"] + assert ("✓ Claude Haiku 4.5", "m:claude-haiku-4-5") in _buttons(keyboard) @pytest.mark.asyncio - async def test_switch_by_id(self): - from memclaw.bot.handlers import model_command_reply + async def test_index_button_press(self): + handlers = _telegram_handlers() + handlers.agent.backend = backend = _SelectableBackend() + await handlers.handle_callback(_callback_update("mi:1"), MagicMock()) + assert backend.model == "claude-haiku-4-5" - backend = _SelectableBackend() - backend.model = "claude-haiku-4-5" - await model_command_reply(backend, ["claude-opus-5"]) + @pytest.mark.asyncio + async def test_stale_button_alerts(self): + handlers = _telegram_handlers() + handlers.agent.backend = backend = _SelectableBackend() + update = _callback_update("m:claude-gone") + await handlers.handle_callback(update, MagicMock()) assert backend.model == "claude-opus-5" + assert update.callback_query.answer.await_args.kwargs["show_alert"] is True + update.callback_query.edit_message_text.assert_not_awaited() @pytest.mark.asyncio - async def test_bad_number_and_unknown_id(self): - from memclaw.bot.handlers import model_command_reply - - backend = _SelectableBackend() - assert "no model number 9" in await model_command_reply(backend, ["9"]) - assert "Unknown model" in await model_command_reply(backend, ["claude-nope"]) + async def test_button_press_from_unknown_user_is_ignored(self): + handlers = _telegram_handlers(allowed=False) + handlers.agent.backend = backend = _SelectableBackend() + update = _callback_update("m:claude-haiku-4-5") + await handlers.handle_callback(update, MagicMock()) assert backend.model == "claude-opus-5" + update.callback_query.edit_message_text.assert_not_awaited() @pytest.mark.asyncio - async def test_unsupported_backend(self): - from memclaw.bot.handlers import effort_command_reply, model_command_reply + async def test_command_with_number_or_id(self): + handlers = _telegram_handlers() + handlers.agent.backend = backend = _SelectableBackend() + update = _command_update() + context = MagicMock() - backend = MagicMock(supports_model_selection=False, display_name="Cursor SDK") - assert "isn't supported for the Cursor SDK backend" in await model_command_reply(backend, []) - assert "isn't supported for the Cursor SDK backend" in await effort_command_reply(backend, ["low"]) + context.args = ["2"] + await handlers.model_command(update, context) + assert backend.model == "claude-haiku-4-5" + + context.args = ["claude-opus-5"] + await handlers.model_command(update, context) + assert backend.model == "claude-opus-5" + assert update.message.reply_text.await_args.kwargs["reply_markup"] is not None @pytest.mark.asyncio - async def test_telegram_handler_passes_args(self): + async def test_bad_number_and_unknown_id(self): handlers = _telegram_handlers() - handlers.agent.backend = _SelectableBackend() + handlers.agent.backend = backend = _SelectableBackend() update = _command_update() context = MagicMock() - context.args = ["2"] - await handlers.model_command(update, context) - assert handlers.agent.backend.model == "claude-haiku-4-5" - update.message.reply_text.assert_awaited_once() + for arg, expected in (("9", "no model number 9"), ("claude-nope", "Unknown model")): + context.args = [arg] + await handlers.model_command(update, context) + assert expected in update.message.reply_text.await_args.args[0] + assert backend.model == "claude-opus-5" + + @pytest.mark.asyncio + async def test_unsupported_backend(self): + from memclaw.bot.handlers import effort_menu, model_menu + + backend = MagicMock(supports_model_selection=False, display_name="Cursor SDK") + text, keyboard = await model_menu(backend) + assert "isn't supported for the Cursor SDK backend" in text + assert keyboard is None + text, keyboard = await effort_menu(backend) + assert "isn't supported for the Cursor SDK backend" in text + assert keyboard is None class TestEffortCommand: @pytest.mark.asyncio - async def test_shows_supported_levels(self): - from memclaw.bot.handlers import effort_command_reply + async def test_menu_has_a_button_per_level(self): + from memclaw.bot.handlers import effort_menu - reply = await effort_command_reply(_SelectableBackend(), []) - assert "Effort: high" in reply - assert "low, medium, high" in reply + text, keyboard = await effort_menu(_SelectableBackend()) + assert "Effort: high" in text + assert _buttons(keyboard) == [ + ("low", "e:low"), ("medium", "e:medium"), ("✓ high", "e:high"), + ] @pytest.mark.asyncio - async def test_model_without_effort(self): - from memclaw.bot.handlers import effort_command_reply + async def test_model_without_effort_has_no_buttons(self): + from memclaw.bot.handlers import effort_menu backend = _SelectableBackend() backend.model, backend.effort = "claude-haiku-4-5", None - reply = await effort_command_reply(backend, []) - assert "doesn't support an effort setting" in reply + text, keyboard = await effort_menu(backend) + assert "doesn't support an effort setting" in text + assert keyboard is None @pytest.mark.asyncio - async def test_sets_level(self): - from memclaw.bot.handlers import effort_command_reply + async def test_button_press_sets_level(self): + handlers = _telegram_handlers() + handlers.agent.backend = backend = _SelectableBackend() + update = _callback_update("e:low") + await handlers.handle_callback(update, MagicMock()) + assert backend.effort == "low" + text = update.callback_query.edit_message_text.await_args.args[0] + assert "Effort set to low." in text - backend = _SelectableBackend() - reply = await effort_command_reply(backend, ["low"]) + @pytest.mark.asyncio + async def test_command_sets_level(self): + handlers = _telegram_handlers() + handlers.agent.backend = backend = _SelectableBackend() + update = _command_update() + context = MagicMock() + context.args = ["low"] + await handlers.effort_command(update, context) assert backend.effort == "low" - assert "Effort set to low." in reply @pytest.mark.asyncio async def test_rejected_level(self): - from memclaw.bot.handlers import effort_command_reply - - backend = _SelectableBackend() - reply = await effort_command_reply(backend, ["max"]) - assert "doesn't support 'max'" in reply + handlers = _telegram_handlers() + handlers.agent.backend = backend = _SelectableBackend() + update = _command_update() + context = MagicMock() + context.args = ["max"] + await handlers.effort_command(update, context) + assert "doesn't support 'max'" in update.message.reply_text.await_args.args[0] assert backend.effort == "high" From 44ae74cf78c3111b2982d4d853499205663891d0 Mon Sep 17 00:00:00 2001 From: Emanuilo Jovanovic Date: Wed, 7 Oct 2026 22:02:42 +0200 Subject: [PATCH 03/10] feat: Telegram progress message for agent tool steps --- README.md | 2 +- memclaw/agent.py | 3 ++ memclaw/backends/base.py | 35 +++++++++++++++- memclaw/backends/claude.py | 8 +++- memclaw/backends/cursor.py | 4 +- memclaw/backends/cursor_sdk_adapter.py | 14 ++++++- memclaw/bot/handlers.py | 35 +++++++++++++++- memclaw/tools.py | 27 +++++++++++++ tests/backends/test_claude.py | 47 +++++++++++++++++++++ tests/backends/test_cursor.py | 31 ++++++++++++++ tests/test_handlers.py | 56 ++++++++++++++++++++++++++ 11 files changed, 256 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index f0a04aa..c4f1412 100644 --- a/README.md +++ b/README.md @@ -74,7 +74,7 @@ Each chat is one ongoing conversation that survives restarts. Send `/new` to sta ### Telegram Bot -The Telegram bot shows a **typing indicator** while processing so you know it's working on your request. +The Telegram bot shows a **typing indicator** while processing, plus a silent "Working… step N" message that follows the agent's tool calls (searching, saving, reminders) and disappears once the reply arrives. Commands: `/new` starts a fresh conversation, `/model` shows the available models as buttons (tap one to switch, or send `/model `), and `/effort` does the same for the reasoning effort. Model and effort changes are saved to `~/.memclaw/.env` and keep the current conversation. They're available with the Claude backend. diff --git a/memclaw/agent.py b/memclaw/agent.py index 2c57df2..3792b59 100644 --- a/memclaw/agent.py +++ b/memclaw/agent.py @@ -24,6 +24,7 @@ from loguru import logger from .backends import AgentBackend, build_backend +from .backends.base import ProgressCallback from .config import MemclawConfig from .index import MemoryIndex from .reminders import ReminderScheduler @@ -403,6 +404,7 @@ async def handle( image_b64: str | None = None, image_media_type: str = "image/jpeg", chat_id: str | None = None, + on_tool: ProgressCallback | None = None, ) -> tuple[str, list[dict]]: self._found_images.clear() self._tools.chat_id = chat_id @@ -449,6 +451,7 @@ async def handle( image_b64=image_b64, image_media_type=image_media_type, max_turns=10, + on_tool=on_tool, ) elapsed_ms = int((time.perf_counter() - t0) * 1000) diff --git a/memclaw/backends/base.py b/memclaw/backends/base.py index 5a28f93..f4046f2 100644 --- a/memclaw/backends/base.py +++ b/memclaw/backends/base.py @@ -18,9 +18,12 @@ class in `memclaw.backends.__init__.REGISTRY`. Nothing else in the project from __future__ import annotations +from collections.abc import Awaitable, Callable from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING, ClassVar, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, ClassVar, Protocol, runtime_checkable + +from loguru import logger if TYPE_CHECKING: from rich.console import Console @@ -52,6 +55,32 @@ class TurnResult: session_id: str | None = None +@dataclass +class ToolStep: + """One tool call the agent makes during a turn, for progress messages.""" + + index: int # 1-based, per turn + name: str # Memclaw tool name, e.g. "memory_search" + summary: str # short human-readable description + + +ProgressCallback = Callable[[ToolStep], Awaitable[None]] + + +async def report_tool_step( + on_tool: ProgressCallback | None, index: int, name: str, args: Any, +) -> None: + """Call *on_tool* for a tool call; a failing callback never fails the turn.""" + if on_tool is None: + return + from ..tools import describe_tool_call # local import — tools imports a lot + + try: + await on_tool(ToolStep(index, name, describe_tool_call(name, args))) + except Exception as exc: + logger.debug("Progress callback failed: {exc}", exc=exc) + + @runtime_checkable class AgentBackend(Protocol): """The contract every agent backend must satisfy.""" @@ -155,6 +184,7 @@ async def run_turn( image_b64: str | None = None, image_media_type: str = "image/jpeg", max_turns: int = 10, + on_tool: ProgressCallback | None = None, ) -> TurnResult: """Run one full agentic turn with tool access, in the conversation identified by *session_key*. @@ -166,6 +196,9 @@ async def run_turn( the session to resume if the backend has to (re)connect; if it can't be resumed the backend starts a fresh one. + ``on_tool`` is called (via `report_tool_step`) for every tool call + the agent makes, for progress messages. + The backend is responsible for translating ``TOOL_DEFINITIONS`` into its SDK's tool format and routing tool calls back through ``tool_executor.execute(name, args)``. diff --git a/memclaw/backends/claude.py b/memclaw/backends/claude.py index 03cdb2c..0a65cd0 100644 --- a/memclaw/backends/claude.py +++ b/memclaw/backends/claude.py @@ -32,7 +32,7 @@ from loguru import logger from ..prompts import choose from ..tools import TOOL_DEFINITIONS -from .base import TurnResult +from .base import ProgressCallback, TurnResult, report_tool_step from .claude_models import SDK_EFFORT_LEVELS, ModelInfo, fetch_models from .mcp_tools import MCP_SERVER_NAME @@ -677,6 +677,7 @@ async def run_turn( image_b64: str | None = None, image_media_type: str = "image/jpeg", max_turns: int = 10, + on_tool: ProgressCallback | None = None, ) -> TurnResult: if self._mcp_server is None: self._mcp_server = _build_mcp_server(tool_executor) @@ -699,6 +700,7 @@ async def run_turn( _compose_prompt(context, user_message), image_b64=image_b64, image_media_type=image_media_type, + on_tool=on_tool, ) except BaseException: # The CLI may be dead or mid-turn; reconnect (resuming) next time. @@ -732,6 +734,7 @@ async def _query( *, image_b64: str | None, image_media_type: str, + on_tool: ProgressCallback | None = None, ) -> TurnResult: """Send one user message on a connected client and collect the reply.""" if image_b64: @@ -741,6 +744,7 @@ async def _query( last_text = "" result = TurnResult(text="") + steps = 0 async for msg in client.receive_response(): if isinstance(msg, AssistantMessage): turn_text = "" @@ -756,6 +760,8 @@ async def _query( if tool_name.startswith(prefix): tool_name = tool_name[len(prefix):] logger.info("Tool call: {name}({args})", name=tool_name, args=args_str) + steps += 1 + await report_tool_step(on_tool, steps, tool_name, block.input) if turn_text: last_text = turn_text elif isinstance(msg, ResultMessage): diff --git a/memclaw/backends/cursor.py b/memclaw/backends/cursor.py index 2a2f996..3e081e3 100644 --- a/memclaw/backends/cursor.py +++ b/memclaw/backends/cursor.py @@ -9,7 +9,7 @@ from loguru import logger -from .base import TurnResult +from .base import ProgressCallback, TurnResult from .cursor_hooks import cursor_hooks_status, ensure_cursor_hooks from .cursor_sdk_adapter import RunUsageTracker, collect_run_result, extract_run_text from .mcp_bridge import HttpMcpServer, mcp_servers_for @@ -324,6 +324,7 @@ async def run_turn( image_b64: str | None = None, image_media_type: str = "image/jpeg", max_turns: int = 10, + on_tool: ProgressCallback | None = None, ) -> TurnResult: from cursor_sdk import CursorAgentError, SendOptions @@ -369,6 +370,7 @@ async def run_turn( run, max_turns=max_turns, usage_tracker=usage_tracker, + on_tool=on_tool, ) finally: await agent.close() diff --git a/memclaw/backends/cursor_sdk_adapter.py b/memclaw/backends/cursor_sdk_adapter.py index ca17810..b36d1a8 100644 --- a/memclaw/backends/cursor_sdk_adapter.py +++ b/memclaw/backends/cursor_sdk_adapter.py @@ -7,7 +7,7 @@ from loguru import logger -from .base import TurnResult +from .base import ProgressCallback, TurnResult, report_tool_step from .mcp_tools import MCP_SERVER_NAME @@ -279,10 +279,12 @@ async def collect_run_result( *, max_turns: int, usage_tracker: RunUsageTracker | None = None, + on_tool: ProgressCallback | None = None, ) -> TurnResult: """Drain the run stream for logging and return a normalized TurnResult.""" last_text = "" pending_tools: dict[str, str] = {} + reported: set[str] = set() async for event in run.events(): message = getattr(event, "sdk_message", None) @@ -290,6 +292,16 @@ async def collect_run_result( turn_text = _handle_sdk_message(message, pending_tools=pending_tools) if turn_text: last_text = turn_text + if getattr(message, "type", None) == "tool_call" and str( + getattr(message, "status", "") + ) == "running": + name, args = normalize_tool_call( + getattr(message, "name", ""), getattr(message, "args", None), + ) + call_id = str(getattr(message, "call_id", "") or len(reported)) + if call_id not in reported: + reported.add(call_id) + await report_tool_step(on_tool, len(reported), name, args) for call_id, name in pending_tools.items(): tool_name, _ = normalize_tool_call(name, {}) diff --git a/memclaw/bot/handlers.py b/memclaw/bot/handlers.py index ed4e8e9..46e0623 100644 --- a/memclaw/bot/handlers.py +++ b/memclaw/bot/handlers.py @@ -8,16 +8,19 @@ import asyncio import base64 +import contextlib +import time from loguru import logger from openai import AsyncOpenAI from telegram import InlineKeyboardButton, InlineKeyboardMarkup, Update from telegram.constants import ChatAction -from telegram.error import BadRequest +from telegram.error import BadRequest, TelegramError from telegram.ext import ContextTypes from ..agent import MemclawAgent from ..backends import AgentBackend +from ..backends.base import ToolStep from ..config import MemclawConfig from ..reminders import ReminderScheduler from . import mask_user_id @@ -44,6 +47,10 @@ async def _typing_loop(bot, chat_id: int): ("effort", "Show or change the reasoning effort"), ] +# Minimum seconds between edits of the "Working…" status message, to stay +# well under Telegram's rate limits. +STATUS_EDIT_INTERVAL_S = 2.0 + # Telegram limits callback_data to 64 bytes. CALLBACK_DATA_MAX_BYTES = 64 @@ -220,15 +227,41 @@ async def _send_with_typing( """Run the agent with a typing indicator, then send the full response.""" chat_id = update.effective_chat.id typing_task = asyncio.create_task(_typing_loop(context.bot, chat_id)) + # One silent status message, edited as the agent works and deleted + # once the reply is ready. + status = None + last_edit = 0.0 + + async def on_tool(step: ToolStep) -> None: + nonlocal status, last_edit + now = time.monotonic() + if now - last_edit < STATUS_EDIT_INTERVAL_S: + return + last_edit = now + text = f"Working… step {step.index}: {step.summary}" + try: + if status is None: + status = await context.bot.send_message( + chat_id, text, disable_notification=True, + ) + else: + await status.edit_text(text) + except TelegramError as exc: + logger.debug("Status update failed: {exc}", exc=exc) + try: response_text, found_images = await self.agent.handle( prompt, image_b64=image_b64, image_media_type=image_media_type, chat_id=str(chat_id), + on_tool=on_tool, ) finally: typing_task.cancel() + if status is not None: + with contextlib.suppress(TelegramError): + await status.delete() await self._send_response(update, context, response_text, found_images) # ------------------------------------------------------------------ diff --git a/memclaw/tools.py b/memclaw/tools.py index 088c4eb..03429fc 100644 --- a/memclaw/tools.py +++ b/memclaw/tools.py @@ -204,6 +204,33 @@ # ── Helpers ────────────────────────────────────────────────────────── +# Progress-message wording per tool: (label, argument shown after it). +_STEP_LABELS: dict[str, tuple[str, str | None]] = { + "memory_save": ("Saving to memory", None), + "memory_search": ("Searching memories", "query"), + "image_save": ("Saving the image", None), + "image_search": ("Searching images", "query"), + "update_instructions": ("Updating my instructions", None), + "file_write": ("Writing", "file_path"), + "file_read": ("Reading", "file_path"), + "reminder_create": ("Setting a reminder", "text"), + "reminder_list": ("Checking reminders", None), + "reminder_cancel": ("Cancelling a reminder", None), +} + + +def describe_tool_call(name: str, args: Any, limit: int = 80) -> str: + """A short one-line description of a tool call for progress messages.""" + label, arg = _STEP_LABELS.get(name, (name, None)) + value = args.get(arg) if arg and isinstance(args, dict) else None + if not isinstance(value, str) or not value.strip(): + return label + value = " ".join(value.split()) + if len(value) > limit: + value = value[: limit - 1] + "…" + return f"{label}: {value}" + + def _format_results(results: list[SearchResult]) -> str: parts = [] for i, r in enumerate(results, 1): diff --git a/tests/backends/test_claude.py b/tests/backends/test_claude.py index 0381b63..b8fad97 100644 --- a/tests/backends/test_claude.py +++ b/tests/backends/test_claude.py @@ -901,3 +901,50 @@ def test_fetch_failure_still_saves_the_credential(self): assert values == {"ANTHROPIC_API_KEY": "sk-ant-new"} assert drops == ["CLAUDE_CODE_OAUTH_TOKEN"] + + +class TestToolSteps: + @pytest.mark.asyncio + async def test_each_tool_call_is_reported(self, claude_config): + from claude_agent_sdk import AssistantMessage, ToolUseBlock + + backend = ClaudeAgentBackend(claude_config(oauth="x")) + fake = _FakeClients("done") + steps = [] + + async def on_tool(step): + steps.append(step) + + with patch("memclaw.backends.claude.ClaudeSDKClient", side_effect=fake): + await _turn(backend, _executor(backend.config)) # builds the client + client = fake.clients[0] + original = client.receive_response.side_effect + + def with_tools(): + async def _gen(): + yield AssistantMessage(content=[ + ToolUseBlock(id="1", name="mcp__memclaw__memory_search", + input={"query": "dog name"}), + ToolUseBlock(id="2", name="mcp__memclaw__memory_save", + input={"content": "x"}), + ], model="claude-sonnet-5-5") + async for msg in original(): + yield msg + return _gen() + + client.receive_response.side_effect = with_tools + await _turn(backend, _executor(backend.config), on_tool=on_tool) + + assert [(s.index, s.name, s.summary) for s in steps] == [ + (1, "memory_search", "Searching memories: dog name"), + (2, "memory_save", "Saving to memory"), + ] + + @pytest.mark.asyncio + async def test_failing_callback_does_not_fail_the_turn(self): + from memclaw.backends.base import report_tool_step + + async def broken(step): + raise RuntimeError("telegram down") + + await report_tool_step(broken, 1, "memory_save", {}) # no exception diff --git a/tests/backends/test_cursor.py b/tests/backends/test_cursor.py index 4744e5d..3f50250 100644 --- a/tests/backends/test_cursor.py +++ b/tests/backends/test_cursor.py @@ -224,6 +224,37 @@ async def _events(): assert "motivational video" in result.text assert result.text != "you." + @pytest.mark.asyncio + async def test_collect_run_result_reports_each_tool_call_once(self): + def _call(call_id, status, name, args): + return SimpleNamespace( + sdk_message=SimpleNamespace( + type="tool_call", status=status, name=name, args=args, call_id=call_id, + ), + ) + + async def _events(): + yield _call("1", "running", "mcp", {"toolName": "memory_search", + "args": {"query": "dog"}}) + yield _call("1", "running", "mcp", {"toolName": "memory_search", + "args": {"query": "dog"}}) + yield _call("1", "completed", "mcp", {"toolName": "memory_search"}) + yield _call("2", "running", "memclaw_memory_save", {"content": "x"}) + + mock_run = AsyncMock() + mock_run.events = MagicMock(return_value=_events()) + mock_run.wait = AsyncMock(return_value=SimpleNamespace(result="ok", num_turns=1)) + steps = [] + + async def on_tool(step): + steps.append(step) + + await collect_run_result(mock_run, max_turns=10, on_tool=on_tool) + assert [(s.index, s.name, s.summary) for s in steps] == [ + (1, "memory_search", "Searching memories: dog"), + (2, "memory_save", "Saving to memory"), + ] + @pytest.mark.asyncio async def test_collect_run_result_reads_turn_ended_usage(self): tracker = RunUsageTracker() diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 27dfaaf..5bf7401 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -482,3 +482,59 @@ async def test_rejected_level(self): await handlers.effort_command(update, context) assert "doesn't support 'max'" in update.message.reply_text.await_args.args[0] assert backend.effort == "high" + + +class TestStepStatusMessage: + @pytest.mark.asyncio + async def test_status_message_is_sent_edited_and_deleted(self): + from memclaw.backends.base import ToolStep + from memclaw.bot import handlers as handlers_mod + + handlers = _telegram_handlers() + handlers._send_response = AsyncMock() + status = MagicMock(edit_text=AsyncMock(), delete=AsyncMock()) + context = MagicMock() + context.bot.send_message = AsyncMock(return_value=status) + context.bot.send_chat_action = AsyncMock() + + async def handle(prompt, *, on_tool, **kwargs): + await on_tool(ToolStep(1, "memory_search", "Searching memories: dog")) + await on_tool(ToolStep(2, "memory_save", "Saving to memory")) + return "Bruno", [] + + handlers.agent.handle = handle + with patch.object(handlers_mod, "STATUS_EDIT_INTERVAL_S", 0): + await handlers._send_with_typing(_command_update(), context, "hi") + + context.bot.send_message.assert_awaited_once() + assert context.bot.send_message.await_args.args[1] == ( + "Working… step 1: Searching memories: dog" + ) + status.edit_text.assert_awaited_once_with("Working… step 2: Saving to memory") + status.delete.assert_awaited_once() + handlers._send_response.assert_awaited_once() + + @pytest.mark.asyncio + async def test_no_status_message_without_tool_calls(self): + handlers = _telegram_handlers() + handlers._send_response = AsyncMock() + context = MagicMock() + context.bot.send_message = AsyncMock() + context.bot.send_chat_action = AsyncMock() + handlers.agent.handle = AsyncMock(return_value=("hello", [])) + + await handlers._send_with_typing(_command_update(), context, "hi") + context.bot.send_message.assert_not_awaited() + + +class TestDescribeToolCall: + def test_labels_and_truncation(self): + from memclaw.tools import describe_tool_call + + assert describe_tool_call("memory_search", {"query": "dog name"}) == ( + "Searching memories: dog name" + ) + assert describe_tool_call("reminder_list", {}) == "Checking reminders" + assert describe_tool_call("unknown_tool", {"x": 1}) == "unknown_tool" + long = describe_tool_call("file_read", {"file_path": "a" * 200}) + assert long.endswith("…") and len(long) <= len("Reading: ") + 80 From 8f28843422afe9e4229796cca14a078cb9ff4767 Mon Sep 17 00:00:00 2001 From: Emanuilo Jovanovic Date: Wed, 7 Oct 2026 21:41:08 +0200 Subject: [PATCH 04/10] feat: native Cursor conversations, /model and /effort for Cursor --- README.md | 6 +- docs/docker.md | 13 +- docs/memory-architecture.md | 13 +- memclaw/backends/cursor.py | 415 +++++++++++++++++++------ memclaw/config.py | 4 +- tests/backends/test_cursor.py | 567 +++++++++++++++++++++++++--------- tests/test_cli_status.py | 4 +- 7 files changed, 752 insertions(+), 270 deletions(-) diff --git a/README.md b/README.md index c4f1412..dca5b42 100644 --- a/README.md +++ b/README.md @@ -76,7 +76,7 @@ Each chat is one ongoing conversation that survives restarts. Send `/new` to sta The Telegram bot shows a **typing indicator** while processing, plus a silent "Working… step N" message that follows the agent's tool calls (searching, saving, reminders) and disappears once the reply arrives. -Commands: `/new` starts a fresh conversation, `/model` shows the available models as buttons (tap one to switch, or send `/model `), and `/effort` does the same for the reasoning effort. Model and effort changes are saved to `~/.memclaw/.env` and keep the current conversation. They're available with the Claude backend. +Commands: `/new` starts a fresh conversation, `/model` shows the available models as buttons (tap one to switch, or send `/model `), and `/effort` does the same for the reasoning effort. Model and effort changes are saved to `~/.memclaw/.env` and keep the current conversation. With Cursor, `/effort` only applies to models that expose a reasoning setting. #### Setup @@ -222,6 +222,7 @@ Each chat keeps one Claude Code process running and reuses it, so replies don't - The wizard prompts for `CURSOR_API_KEY` and an optional `CURSOR_MODEL` (defaults to `composer-2.5`). - Memclaw tools are exposed to the Cursor agent via a long-lived local HTTP MCP server on `127.0.0.1` (default port `17373`, override with `MEMCLAW_MCP_PORT`). - Requests bill per token against your Cursor plan (or pay-as-you-go overage rates once your plan limit is reached), and the cost is logged with each turn. +- One Cursor SDK bridge process serves all chats, and each chat keeps one local Cursor agent open across turns. The agent stores its conversation under `~/.cursor` and is resumed after a restart. ## How It Works @@ -362,7 +363,8 @@ Backend choice and credentials are covered above in [Agent Backend](#agent-backe | `AGENT_BACKEND` | Optional | Agent SDK to use (defaults to `claude`; set to `cursor` for Cursor SDK) | | `MEMCLAW_PLATFORM` | Optional | Front-end the bare `memclaw` launches: `telegram`, `whatsapp`, `slack`, or `terminal` (defaults to `terminal`) | | `CURSOR_API_KEY` | For Cursor backend | Cursor API key from Dashboard → Integrations | -| `CURSOR_MODEL` | For Cursor backend | Cursor model name (defaults to `composer-2.5`) | +| `CURSOR_MODEL` | For Cursor backend | Cursor model name (defaults to `composer-2.5`; also set by `/model`) | +| `CURSOR_EFFORT` | Optional | Cursor reasoning level, for models that have one (set by `/effort`) | | `MEMCLAW_MCP_PORT` | For Cursor backend | Local MCP HTTP port (defaults to `17373`) | | `TELEGRAM_BOT_TOKEN` | For Telegram bot | Your Telegram bot token | | `ALLOWED_USER_IDS` | For Telegram bot | Comma-separated Telegram user IDs | diff --git a/docs/docker.md b/docs/docker.md index 2bd5b8c..5d1fc8a 100644 --- a/docs/docker.md +++ b/docs/docker.md @@ -123,13 +123,16 @@ docker compose up -d --build grows with your memory vault, so back the volume up periodically. - The Claude backend (`AGENT_BACKEND=claude`, the default) keeps one Claude Code CLI process running per active chat (closed on `/new` - and shutdown). The Cursor backend (`AGENT_BACKEND=cursor`) is - pure-Python and slightly leaner if you don't need Claude. + and shutdown). The Cursor backend (`AGENT_BACKEND=cursor`) runs a + single Cursor SDK bridge process (Node) shared by all chats. - Conversations resume after a restart: `sessions.json` lives in the data volume and Claude Code's own session files in the container's `~/.claude`. That directory isn't on the volume, so after a rebuild (`docker compose up -d --build` recreates the container) each chat - starts a fresh conversation. Memories are unaffected. + starts a fresh conversation. Memories are unaffected. The same goes + for the Cursor backend, whose agents are stored in the container's + `~/.cursor`. - Telegram's `/model` and `/effort` save their choice to - `~/.memclaw/.env`. A `CLAUDE_MODEL` / `CLAUDE_EFFORT` passed in from - the compose `.env` wins over that file on the next start. + `~/.memclaw/.env`. A `CLAUDE_MODEL` / `CLAUDE_EFFORT` (or + `CURSOR_MODEL` / `CURSOR_EFFORT`) passed in from the compose `.env` + wins over that file on the next start. diff --git a/docs/memory-architecture.md b/docs/memory-architecture.md index 575d8c9..2ec913b 100644 --- a/docs/memory-architecture.md +++ b/docs/memory-architecture.md @@ -10,17 +10,18 @@ Short-term memory is the conversation itself: what you and the agent said in the ### Conversation Sessions -Each chat's conversation lives in the agent backend's own session (for Claude, a Claude Code session), not in a buffer Memclaw re-sends with every message. +Each chat's conversation lives in the agent backend's own session (for Claude, a Claude Code session; for Cursor, a local Cursor SDK agent), not in a buffer Memclaw re-sends with every message. **How it works:** - Every chat gets a session key: `:` for the bots (e.g. `telegram:12345`, `slack:C01ABC`), `cli` for the interactive terminal. - The Claude backend keeps one connected Claude CLI process per chat and reuses it for every turn, so a message doesn't pay the CLI start-up cost. Turns in the same chat are serialized. +- The Cursor backend runs one Cursor SDK bridge process for all chats and keeps one open agent per chat. The agent stores its conversation on disk (under `~/.cursor`, in a store tied to the bridge's working directory, which is pinned to the memory directory). Cursor agents have no system-prompt field, so the system prompt goes at the top of a new agent's first message; later messages carry only the context block and the message. - The session id is saved in `~/.memclaw/sessions.json` (written atomically). After a restart, the next message in that chat resumes the same session, so the conversation carries on. - If a session can't be resumed (deleted, created on another machine, ...), a warning is logged and a fresh session starts. -- Each saved session records a **fingerprint** of the code-level instructions (the system-prompt template and the backend). Claude Code keeps a session's original system prompt when it resumes, so a session created under different instructions (e.g. after an upgrade that changed the template) is not resumed; the chat starts fresh. +- Each saved session records a **fingerprint** of the code-level instructions (the system-prompt template and the backend). A resumed session keeps its original system prompt, so a session created under different instructions (e.g. after an upgrade that changed the template) is not resumed; the chat starts fresh. - `/new` (Telegram command; a message of exactly `/new` in Slack, WhatsApp and the terminal) closes that chat's session and forgets its id. Your memories are untouched. -- `/model` and `/effort` (Telegram) switch the model or effort level. Live sessions reconnect on their next message, resuming the same session with the new settings, so the conversation is kept. +- `/model` and `/effort` (Telegram) switch the model or effort level, and the conversation is kept. Claude sessions reconnect on their next message, resuming the same session with the new settings; Cursor passes the model with every message, so its agents just use it. For Cursor, `/effort` is available only on models that expose a reasoning/effort parameter. **What this means in practice:** @@ -31,9 +32,7 @@ You: What's my dog called? Assistant: Your dog's name is Max. ``` -The conversation is bounded only by the backend's own context management (Claude Code compacts long sessions on its own). Use `/new` to start over. - -The Cursor backend has no native sessions yet: as an interim measure it keeps the last 10 messages per chat in memory and folds them into each prompt. That transcript does not survive a restart. +The conversation is bounded only by the backend's own context management (e.g. Claude Code compacts long sessions on its own). Use `/new` to start over. ### System Prompt and Per-Message Context @@ -56,7 +55,7 @@ The block stays in the session transcript, which is why it is kept small. The session stores a hash of the AGENTS.md and MEMORY.md content it has seen (persisted in `sessions.json`). Before each message the files are hashed again; if either changed (an `update_instructions` call, a consolidation, a manual edit), its new content is included **once** in that message's context block, labelled as superseding the earlier version, and the stored hash is updated. The conversation continues without a reset. Changes the session makes itself during a turn (`memory_save` with `permanent=true`, `update_instructions`) count as seen and aren't sent back; a consolidation that rewrites MEMORY.md during a turn is. -Because Claude Code keeps a resumed session's original system prompt, the stored hashes always describe what the session has actually seen, including across restarts. +Because a resumed session keeps its original system prompt, the stored hashes always describe what the session has actually seen, including across restarts. ## Long-Term Memory diff --git a/memclaw/backends/cursor.py b/memclaw/backends/cursor.py index 3e081e3..9c2da99 100644 --- a/memclaw/backends/cursor.py +++ b/memclaw/backends/cursor.py @@ -1,15 +1,25 @@ -"""Cursor Python SDK agent backend.""" +"""Cursor Python SDK agent backend. + +One `cursor-sdk-bridge` process (a Node subprocess) serves the whole Memclaw +process. Each chat keeps one open local agent on it, reused across turns: the +agent stores its own conversation on disk, so follow-up messages only carry +the per-turn context and the new message, and a restarted process resumes the +agent by id. +""" from __future__ import annotations +import asyncio import os -from collections import deque +from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar +import httpx from loguru import logger from .base import ProgressCallback, TurnResult +from .claude_models import ModelInfo from .cursor_hooks import cursor_hooks_status, ensure_cursor_hooks from .cursor_sdk_adapter import RunUsageTracker, collect_run_result, extract_run_text from .mcp_bridge import HttpMcpServer, mcp_servers_for @@ -30,11 +40,10 @@ # Scrub Claude credentials when switching to Cursor so they can't shadow selection. _DROP_KEYS = ["ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN"] -# INTERIM: the Cursor backend has no native conversation sessions yet, so it -# keeps the last few messages per session key in memory and folds them into -# each prompt. To be replaced by native Cursor conversation handling. -_TRANSCRIPT_MESSAGES = 10 -_TRANSCRIPT_MESSAGE_CHARS = 1000 +# A model parameter whose id contains one of these is treated as its effort +# setting (/effort). Cursor reports parameters per model, so there is no +# fixed list of levels. +_EFFORT_PARAM_HINTS = ("effort", "reasoning") def _cursor_api_key(config: "MemclawConfig") -> str: @@ -46,6 +55,11 @@ def _cursor_model(config: "MemclawConfig") -> str: return model or _DEFAULT_MODEL +def _cursor_effort(config: "MemclawConfig") -> str | None: + effort = (config.cursor_effort or os.environ.get("CURSOR_EFFORT", "")).strip() + return effort or None + + def _apply_cost_fallback(result: TurnResult, *, bills_per_token: bool) -> None: if result.cost_usd is not None or not bills_per_token: return @@ -71,31 +85,23 @@ def _build_combined_prompt(*, system_prompt: str, user_message: str) -> str: ) -def _format_transcript(transcript: list[tuple[str, str]]) -> str: - """INTERIM: render recent (role, text) messages for the prompt.""" - if not transcript: - return "" - lines = ["=== RECENT CONVERSATION ==="] - for role, text in transcript: - lines.append(f"{'User' if role == 'user' else 'Assistant'}: {text}") - return "\n".join(lines) - - def _build_user_message( *, system_prompt: str, user_message: str, context: str = "", - transcript: list[tuple[str, str]] | None = None, image_b64: str | None = None, image_media_type: str = "image/jpeg", ) -> str | Any: + """The message to send. Cursor agents have no system prompt field, so a + new agent gets it at the top of its first message; follow-ups pass "".""" from cursor_sdk import SDKImage, UserMessage - parts = [_format_transcript(transcript or []), context, user_message] - prompt = _build_combined_prompt( - system_prompt=system_prompt, - user_message="\n\n".join(p for p in parts if p), + body = "\n\n".join(p for p in (context, user_message) if p) + prompt = ( + _build_combined_prompt(system_prompt=system_prompt, user_message=body) + if system_prompt + else body ) if not image_b64: return prompt @@ -119,7 +125,7 @@ def _agent_options( *, api_key: str, cwd: str, - model: str, + model: Any, mcp_servers: dict[str, Any] | None = None, ) -> Any: from cursor_sdk import AgentOptions @@ -132,23 +138,60 @@ def _agent_options( ) +def _effort_param(model: Any) -> Any | None: + """The parameter of an SDKModel that acts as its effort level, if any.""" + for param in model.parameters: + if param.values and any(h in param.id.lower() for h in _EFFORT_PARAM_HINTS): + return param + return None + + +async def _drain_stderr(stream: asyncio.StreamReader) -> None: + """Keep reading the bridge's stderr. The SDK only reads it until the + bridge is up, and a long-lived bridge would block once the pipe fills.""" + while line := await stream.readline(): + logger.debug("cursor-sdk-bridge: {line}", line=line.decode(errors="replace").rstrip()) + + +def _bridge_failed(exc: BaseException) -> bool: + """True when *exc* came from not reaching the bridge process at all + (it crashed or hung), rather than from the agent or Cursor's API.""" + return isinstance(exc.__cause__, httpx.RequestError) + + +@dataclass +class _Session: + """One chat's open Cursor agent.""" + + # Serializes turns (and /new) within one chat. + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + agent: Any = None + # The bridge client *agent* lives on; a relaunched bridge doesn't know it. + client: Any = None + + class CursorAgentBackend: """Cursor Python SDK implementation of the AgentBackend protocol.""" name: ClassVar[str] = "cursor" display_name: ClassVar[str] = "Cursor SDK" - supports_model_selection: ClassVar[bool] = False + supports_model_selection: ClassVar[bool] = True def __init__(self, config: "MemclawConfig") -> None: self.config = config self._api_key = _cursor_api_key(config) self._model = _cursor_model(config) + self._effort = _cursor_effort(config) self._cwd = str(config.memory_dir) os.environ["MEMCLAW_MEMORY_DIR"] = self._cwd self.bills_per_token = True self._mcp_server = HttpMcpServer() - # INTERIM: session key -> recent (role, text) messages. - self._transcripts: dict[str, deque[tuple[str, str]]] = {} + self._client: Any = None # the shared bridge client, launched lazily + self._client_lock = asyncio.Lock() + self._background: set[asyncio.Task] = set() # stderr drains + self._sessions: dict[str, _Session] = {} + self._models: list[ModelInfo] = [] # filled by list_models() + self._effort_params: dict[str, str] = {} # model id -> effort param id @classmethod def is_configured(cls, config: "MemclawConfig") -> bool: @@ -169,7 +212,10 @@ def configuration_help(cls) -> str: @classmethod def status_rows(cls, config: "MemclawConfig") -> list[tuple[str, str]]: - return [("Model", _cursor_model(config))] + return [ + ("Model", _cursor_model(config)), + ("Effort", _cursor_effort(config) or "default"), + ] @classmethod def wizard_setup( @@ -189,6 +235,7 @@ def wizard_setup( raise SystemExit(1) values: dict[str, str] = {"CURSOR_API_KEY": value} + drop_keys = list(_DROP_KEYS) model_current = existing.get("CURSOR_MODEL", os.environ.get("CURSOR_MODEL", "")) from rich.prompt import Prompt @@ -200,6 +247,9 @@ def wizard_setup( ) if model_answer.strip(): values["CURSOR_MODEL"] = model_answer.strip() + if model_answer.strip() != model_current: + # An effort level picked with /effort belongs to the old model. + drop_keys.append("CURSOR_EFFORT") elif model_current: values["CURSOR_MODEL"] = model_current @@ -214,7 +264,9 @@ def wizard_setup( status=status, ) - return values, list(_DROP_KEYS) + return values, drop_keys + + # -- Lifecycle --------------------------------------------------------- async def on_agent_start(self, tool_executor: "ToolExecutor") -> None: await self._ensure_mcp_server(tool_executor) @@ -230,6 +282,11 @@ async def on_agent_start(self, tool_executor: "ToolExecutor") -> None: ) async def on_agent_shutdown(self) -> None: + for session in self._sessions.values(): + await self._close_agent(session) + self._sessions.clear() + if self._client is not None: + await self._drop_client(self._client) await self._mcp_server.stop() async def _ensure_mcp_server(self, tool_executor: "ToolExecutor") -> None: @@ -238,17 +295,91 @@ async def _ensure_mcp_server(self, tool_executor: "ToolExecutor") -> None: port=self.config.mcp_http_port, ) - async def reset_session(self, session_key: str) -> None: - self._transcripts.pop(session_key, None) + # -- Bridge ------------------------------------------------------------ - def _remember(self, session_key: str, role: str, text: str) -> None: - transcript = self._transcripts.setdefault( - session_key, deque(maxlen=_TRANSCRIPT_MESSAGES), + async def _launch_client(self) -> Any: + from cursor_sdk import AsyncClient + from cursor_sdk._vendor import resolve_bridge_path + + # Local agents are stored under a path derived from the bridge + # process's working directory (not LocalAgentOptions.cwd), and + # launch_bridge can't set it, so start the bridge through `sh` in + # the memory dir. A restarted process then finds its agents again. + Path(self._cwd).mkdir(parents=True, exist_ok=True) + client = await AsyncClient.launch_bridge( + ["/bin/sh", "-c", 'cd "$0" && exec "$@"', self._cwd, resolve_bridge_path()], + workspace=self._cwd, ) - transcript.append((role, text[:_TRANSCRIPT_MESSAGE_CHARS])) + process = getattr(client._owned_bridge, "process", None) + if process is not None and process.stderr is not None: + task = asyncio.create_task(_drain_stderr(process.stderr)) + self._background.add(task) + task.add_done_callback(self._background.discard) + return client + + async def _get_client(self) -> Any: + async with self._client_lock: + if self._client is None: + self._client = await self._launch_client() + return self._client + + async def _drop_client(self, client: Any) -> None: + """Shut *client*'s bridge down; the next call launches a new one.""" + if self._client is client: + self._client = None + try: + await client.aclose() + except Exception as exc: + logger.warning("Error while closing the Cursor SDK bridge: {exc}", exc=exc) - # Model selection isn't supported for Cursor yet (see - # supports_model_selection); CURSOR_MODEL picks the model. + # -- Sessions ---------------------------------------------------------- + + @staticmethod + async def _close_agent(session: _Session) -> None: + """Release the agent's executor. Its conversation stays on disk.""" + agent, session.agent, session.client = session.agent, None, None + if agent is None: + return + try: + await agent.close() + except Exception as exc: + logger.warning("Error while closing Cursor agent: {exc}", exc=exc) + + async def _open_agent( + self, client: Any, *, resume: str | None, mcp_servers: dict[str, Any], + ) -> tuple[Any, bool]: + """Resume agent *resume*, or create a new one when there is none or it + can't be resumed (deleted, from another machine, ...). Returns the + agent and whether it is new.""" + from cursor_sdk import CursorAgentError + + # Local resume doesn't restore the model; pass it along explicitly. + options = _agent_options( + api_key=self._api_key, + cwd=self._cwd, + model=await self._model_selection(), + mcp_servers=mcp_servers, + ) + if resume: + try: + return await client.agents.resume(resume, options), False + except CursorAgentError as exc: + if _bridge_failed(exc): + raise + logger.warning( + "Could not resume Cursor agent {id} ({exc}); starting a fresh one", + id=resume, exc=exc.message, + ) + return await client.agents.create(options), True + + async def reset_session(self, session_key: str) -> None: + session = self._sessions.get(session_key) + if session is None: + return + async with session.lock: + await self._close_agent(session) + + # -- Model selection --------------------------------------------------- @property def model(self) -> str: @@ -256,28 +387,109 @@ def model(self) -> str: @property def effort(self) -> str | None: - return None - - async def list_models(self) -> list[Any]: - raise NotImplementedError("Model selection isn't supported for Cursor yet.") + return self._effort + + async def list_models(self) -> list[ModelInfo]: + client = await self._get_client() + models = await client.models.list(api_key=self._api_key) + infos: list[ModelInfo] = [] + params: dict[str, str] = {} + for m in models: + param = _effort_param(m) + if param is not None: + params[m.id] = param.id + infos.append(ModelInfo( + id=m.id, + display_name=m.display_name or m.id, + created_at="", + effort_levels=[v.value for v in param.values] if param else [], + )) + self._models, self._effort_params = infos, params + return infos + + async def _model_info(self, model_id: str) -> ModelInfo | None: + """*model_id*'s entry in the model list, or None if it isn't listed + or the list can't be fetched.""" + if not self._models: + try: + await self.list_models() + except Exception as exc: + logger.warning("Could not fetch the Cursor model list: {exc}", exc=exc) + return None + return next((m for m in self._models if m.id == model_id), None) + + async def _model_selection(self) -> Any: + """The model (and effort parameter, if set) to run with.""" + from cursor_sdk import ModelParameterValue, ModelSelection + + if not self._effort: + return ModelSelection(id=self._model) + if self._model not in self._effort_params: + await self._model_info(self._model) # learn the parameter's id + param = self._effort_params.get(self._model) + if param is None: + return ModelSelection(id=self._model) + return ModelSelection( + id=self._model, + params=[ModelParameterValue(id=param, value=self._effort)], + ) async def effort_levels(self) -> list[str] | None: - return None + info = await self._model_info(self._model) + return list(info.effort_levels) if info else None async def set_model(self, model_id: str) -> str: - raise NotImplementedError("Model selection isn't supported for Cursor yet.") + model_id = model_id.strip() + info = await self._model_info(model_id) + if info is None and self._models: + raise ValueError(f"Unknown model '{model_id}'.") + effort = self._effort + note = "" + if info is not None and self._effort and self._effort not in info.effort_levels: + effort = None + note = ( + f"{info.display_name} doesn't support '{self._effort}' effort, " + "so its default is used." + ) + self._apply_model(model_id, effort) + return note async def set_effort(self, level: str) -> None: - raise NotImplementedError("Effort selection isn't supported for Cursor yet.") + level = level.strip() + levels = await self.effort_levels() + if levels is None: + raise ValueError("Couldn't load the model list right now; try again later.") + if not levels: + raise ValueError("The current model doesn't support an effort setting.") + match = next((v for v in levels if v.lower() == level.lower()), None) + if match is None: + raise ValueError( + f"The current model doesn't support '{level}'. " + f"Supported: {', '.join(levels)}." + ) + self._apply_model(self._model, match) - async def _launch_client(self) -> Any: - from cursor_sdk import AsyncClient + def _apply_model(self, model: str, effort: str | None) -> None: + """Use *model* / *effort* from the next turn on and save them to + ~/.memclaw/.env. The model is passed with every send, so open agents + keep their conversation and need no reconnect.""" + from ..setup import update_env_file # local import — setup imports backends - return await AsyncClient.launch_bridge(workspace=self._cwd) + logger.info( + "Cursor model: {m} effort: {e} -> {m2} effort: {e2}", + m=self._model, e=self._effort, m2=model, e2=effort, + ) + self._model, self._effort = model, effort + self.config.cursor_model = model + self.config.cursor_effort = effort or "" + if effort: + update_env_file({"CURSOR_MODEL": model, "CURSOR_EFFORT": effort}) + else: + update_env_file({"CURSOR_MODEL": model}, drop_keys=["CURSOR_EFFORT"]) - async def run_one_shot(self, *, system_prompt: str, user_message: str) -> str: - from contextlib import asynccontextmanager + # -- Runtime ----------------------------------------------------------- + async def run_one_shot(self, *, system_prompt: str, user_message: str) -> str: from cursor_sdk import AsyncAgent, CursorAgentError if not self._api_key: @@ -290,27 +502,22 @@ async def run_one_shot(self, *, system_prompt: str, user_message: str) -> str: options = _agent_options( api_key=self._api_key, cwd=self._cwd, - model=self._model, + model=await self._model_selection(), ) - @asynccontextmanager - async def client_ctx(): - client = await self._launch_client() - try: - yield client - finally: - await client.aclose() - + client = await self._get_client() try: - async with client_ctx() as client: - result = await AsyncAgent.prompt(prompt, options, client=client) - text = extract_run_text(result) - if not text.strip(): - return "I couldn't generate a response." - return text + # A throwaway agent on the shared bridge (created, run, closed). + result = await AsyncAgent.prompt(prompt, options, client=client) except CursorAgentError as exc: + if _bridge_failed(exc): + await self._drop_client(client) logger.error("Cursor SDK one-shot failed: {msg}", msg=exc.message) raise RuntimeError(f"Cursor SDK error: {exc.message}") from exc + text = extract_run_text(result) + if not text.strip(): + return "I couldn't generate a response." + return text async def run_turn( self, @@ -326,44 +533,49 @@ async def run_turn( max_turns: int = 10, on_tool: ProgressCallback | None = None, ) -> TurnResult: - from cursor_sdk import CursorAgentError, SendOptions + from cursor_sdk import CursorAgentError, LocalSendOptions, SendOptions if not self._api_key: raise RuntimeError("CURSOR_API_KEY is not configured") await self._ensure_mcp_server(tool_executor) - - message = _build_user_message( - system_prompt=system_prompt, - context=context, - user_message=user_message, - transcript=list(self._transcripts.get(session_key, ())), - image_b64=image_b64, - image_media_type=image_media_type, - ) - mcp_config = self._mcp_server.config if mcp_config is None: raise RuntimeError("MCP server failed to start") - mcp_servers = mcp_servers_for(mcp_config) - usage_tracker = RunUsageTracker() - client = await self._launch_client() - try: - agent = await client.agents.create( - _agent_options( - api_key=self._api_key, - cwd=self._cwd, - model=self._model, - mcp_servers=mcp_servers, - ) - ) + + session = self._sessions.setdefault(session_key, _Session()) + async with session.lock: + client = await self._get_client() + if session.agent is not None and session.client is not client: + session.agent = None # its bridge was relaunched + opened = is_new = False try: - run = await agent.send( + if session.agent is None: + session.agent, is_new = await self._open_agent( + client, resume=resume_session_id, mcp_servers=mcp_servers, + ) + session.client = client + opened = True + message = _build_user_message( + system_prompt=system_prompt if is_new else "", + context=context, + user_message=user_message, + image_b64=image_b64, + image_media_type=image_media_type, + ) + usage_tracker = RunUsageTracker() + run = await session.agent.send( message, SendOptions( + # Passed on every send: a resumed agent doesn't + # remember its model, and /model takes effect here. + model=await self._model_selection(), mcp_servers=mcp_servers, on_delta=usage_tracker.on_delta, + # A resumed agent may still show a run that died with + # the previous process; take it over. + local=LocalSendOptions(force=True) if opened and not is_new else None, ), ) result = await collect_run_result( @@ -372,17 +584,18 @@ async def run_turn( usage_tracker=usage_tracker, on_tool=on_tool, ) - finally: - await agent.close() - - if not result.text.strip(): - result.text = "I couldn't generate a response." - _apply_cost_fallback(result, bills_per_token=self.bills_per_token) - self._remember(session_key, "user", user_message) - self._remember(session_key, "assistant", result.text) - return result - except CursorAgentError as exc: - logger.error("Cursor SDK turn failed: {msg}", msg=exc.message) - raise RuntimeError(f"Cursor SDK error: {exc.message}") from exc - finally: - await client.aclose() + result.session_id = session.agent.agent_id + except BaseException as exc: + # The agent may be mid-run or gone; reopen (resuming) next time. + await self._close_agent(session) + if _bridge_failed(exc): + await self._drop_client(client) + if isinstance(exc, CursorAgentError): + logger.error("Cursor SDK turn failed: {msg}", msg=exc.message) + raise RuntimeError(f"Cursor SDK error: {exc.message}") from exc + raise + + if not result.text.strip(): + result.text = "I couldn't generate a response." + _apply_cost_fallback(result, bills_per_token=self.bills_per_token) + return result diff --git a/memclaw/config.py b/memclaw/config.py index 97ca69d..f7da71b 100644 --- a/memclaw/config.py +++ b/memclaw/config.py @@ -2,7 +2,6 @@ import importlib.resources import os -import shutil from dataclasses import dataclass, field from datetime import date from pathlib import Path @@ -47,6 +46,7 @@ class MemclawConfig: # Cursor SDK backend settings cursor_api_key: str = "" cursor_model: str = "" + cursor_effort: str = "" mcp_http_port: int = 17373 # Memory consolidation @@ -81,6 +81,8 @@ def __post_init__(self): self.cursor_api_key = os.environ.get("CURSOR_API_KEY", "") if not self.cursor_model: self.cursor_model = os.environ.get("CURSOR_MODEL", "") + if not self.cursor_effort: + self.cursor_effort = os.environ.get("CURSOR_EFFORT", "") env_mcp_port = os.environ.get("MEMCLAW_MCP_PORT", "").strip() if env_mcp_port: try: diff --git a/tests/backends/test_cursor.py b/tests/backends/test_cursor.py index 3f50250..a21f547 100644 --- a/tests/backends/test_cursor.py +++ b/tests/backends/test_cursor.py @@ -3,12 +3,18 @@ from __future__ import annotations import os +from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch import pytest -from cursor_sdk import HttpMcpServerConfig +from cursor_sdk import ( + HttpMcpServerConfig, + ModelParameterDefinition, + ModelParameterDefinitionValue, + SDKModel, +) from memclaw.backends.base import TurnResult from memclaw.backends import REGISTRY, get_backend_class @@ -35,7 +41,7 @@ @pytest.fixture(autouse=True) def _isolate_credentials(monkeypatch): """Prevent shell env from leaking into MemclawConfig.""" - for name in ("CURSOR_API_KEY", "CURSOR_MODEL"): + for name in ("CURSOR_API_KEY", "CURSOR_MODEL", "CURSOR_EFFORT"): monkeypatch.delenv(name, raising=False) @@ -44,6 +50,7 @@ def _make_config( *, cursor_api_key: str = "", cursor_model: str = "", + cursor_effort: str = "", ) -> MemclawConfig: return MemclawConfig( memory_dir=tmp_path / "m", @@ -51,6 +58,7 @@ def _make_config( anthropic_api_key="test-anthropic-key", cursor_api_key=cursor_api_key, cursor_model=cursor_model, + cursor_effort=cursor_effort, ) @@ -98,11 +106,17 @@ def test_init_reads_config(self, tmp_path): def test_status_rows_show_the_model(self, tmp_path): cfg = _make_config(tmp_path, cursor_model="composer-3") - assert CursorAgentBackend.status_rows(cfg) == [("Model", "composer-3")] + assert CursorAgentBackend.status_rows(cfg) == [ + ("Model", "composer-3"), ("Effort", "default"), + ] + + def test_status_rows_show_the_effort(self, tmp_path): + cfg = _make_config(tmp_path, cursor_effort="high") + assert ("Effort", "high") in CursorAgentBackend.status_rows(cfg) def test_status_rows_fall_back_to_default_model(self, tmp_path): rows = CursorAgentBackend.status_rows(_make_config(tmp_path)) - assert rows == [("Model", "composer-2.5")] + assert rows[0] == ("Model", "composer-2.5") class TestPromptBuilding: @@ -140,6 +154,12 @@ def test_combined_prompt_includes_system_and_user(self): assert "Hello" in prompt assert "---" in prompt + def test_follow_up_message_has_no_system_prompt(self): + message = _build_user_message( + system_prompt="", context="c", user_message="Hi", + ) + assert message == "c\n\nHi" + def test_image_uses_sdk_image(self): message = _build_user_message( system_prompt="Sys", @@ -296,95 +316,288 @@ def test_record_interaction_usage_falls_back_to_token_delta(self): assert result.output_tokens == 42 +_MCP = HttpMcpServerConfig(url="http://127.0.0.1:8765/mcp", type="http") + + +def _make_run(reply: str = "Turn response", num_turns: int = 1): + run = AsyncMock() + run.events = MagicMock(return_value=_empty_events()) + run.wait = AsyncMock(return_value=SimpleNamespace(result=reply, num_turns=num_turns)) + return run + + +def _make_agent(agent_id: str = "agent-1", replies=None): + """A mock AsyncAgent that records each sent message.""" + agent = MagicMock() + agent.agent_id = agent_id + agent.sent = [] + replies = iter(replies or ["Turn response"] * 10) + + async def _send(message, options): + agent.sent.append((message, options)) + return _make_run(next(replies)) + + agent.send = AsyncMock(side_effect=_send) + agent.close = AsyncMock() + return agent + + +def _make_client(*, create=None, resume=None): + client = MagicMock() + client.agents.create = AsyncMock(side_effect=create or [_make_agent()]) + client.agents.resume = AsyncMock(side_effect=resume or []) + client.models.list = AsyncMock(return_value=[]) + client.aclose = AsyncMock() + return client + + +@contextmanager +def _patched(backend, *clients): + """Run *backend* against mock bridge clients (one per launch) and a stub MCP server.""" + with patch.object( + backend, "_launch_client", new_callable=AsyncMock, side_effect=list(clients), + ) as launch, patch.object( + backend, "_ensure_mcp_server", new_callable=AsyncMock + ), patch.object( + HttpMcpServer, "config", new_callable=PropertyMock, return_value=_MCP + ): + yield launch + + +async def _turn(backend, text="User", *, key="cli", resume=None, context="ctx"): + return await backend.run_turn( + system_prompt="System", context=context, user_message=text, + tool_executor=MagicMock(), session_key=key, resume_session_id=resume, + ) + + +def _backend(tmp_path, **kwargs) -> CursorAgentBackend: + return CursorAgentBackend(_make_config(tmp_path, cursor_api_key="crsr_test_key", **kwargs)) + + +def _sdk_model(model_id, *, effort_values=()): + params = () + if effort_values: + params = (ModelParameterDefinition( + id="reasoning", + values=tuple(ModelParameterDefinitionValue(value=v) for v in effort_values), + ),) + return SDKModel(id=model_id, display_name=model_id.title(), parameters=params) + + class TestCursorAgentBackendRuns: @pytest.mark.asyncio async def test_run_one_shot(self, tmp_path): - cfg = _make_config(tmp_path, cursor_api_key="crsr_test_key") - backend = CursorAgentBackend(cfg) - - mock_result = SimpleNamespace(result="Hello from Cursor") + backend = _backend(tmp_path) + client = _make_client() - with patch("cursor_sdk.AsyncClient") as mock_client_cls, patch( + with _patched(backend, client), patch( "cursor_sdk.AsyncAgent.prompt", new_callable=AsyncMock ) as mock_prompt: - mock_client = AsyncMock() - mock_client_cls.launch_bridge = AsyncMock(return_value=mock_client) - mock_prompt.return_value = mock_result - - text = await backend.run_one_shot( - system_prompt="System", - user_message="User", - ) + mock_prompt.return_value = SimpleNamespace(result="Hello from Cursor") + text = await backend.run_one_shot(system_prompt="System", user_message="User") assert text == "Hello from Cursor" mock_prompt.assert_awaited_once() call_prompt = mock_prompt.await_args.args[0] assert "System" in call_prompt assert "User" in call_prompt + assert mock_prompt.await_args.kwargs["client"] is client @pytest.mark.asyncio async def test_run_one_shot_missing_key_raises(self, tmp_path): - cfg = _make_config(tmp_path) - backend = CursorAgentBackend(cfg) - + backend = CursorAgentBackend(_make_config(tmp_path)) with pytest.raises(RuntimeError, match="CURSOR_API_KEY"): await backend.run_one_shot(system_prompt="S", user_message="U") + @pytest.mark.asyncio + async def test_launch_runs_bridge_in_memory_dir(self, tmp_path): + """The agent store path follows the bridge's cwd, so it must be fixed.""" + backend = _backend(tmp_path) + with patch( + "cursor_sdk.AsyncClient.launch_bridge", new_callable=AsyncMock, + return_value=MagicMock(_owned_bridge=None), + ) as launch: + await backend._launch_client() + argv = launch.await_args.args[0] + assert argv[:2] == ["/bin/sh", "-c"] + assert argv[3] == backend._cwd + assert launch.await_args.kwargs["workspace"] == backend._cwd + @pytest.mark.asyncio async def test_run_turn(self, tmp_path): - cfg = _make_config(tmp_path, cursor_api_key="crsr_test_key") - backend = CursorAgentBackend(cfg) + backend = _backend(tmp_path) + agent = _make_agent(replies=["Turn response"]) + client = _make_client(create=[agent]) - mock_run = AsyncMock() - mock_run.events = MagicMock(return_value=_empty_events()) - mock_run.wait = AsyncMock(return_value=SimpleNamespace(result="Turn response", num_turns=2)) - mock_agent = AsyncMock() - mock_agent.send = AsyncMock(return_value=mock_run) - mock_agent.close = AsyncMock() - mock_client = AsyncMock() - mock_client.agents.create = AsyncMock(return_value=mock_agent) - mock_mcp = HttpMcpServerConfig(url="http://127.0.0.1:8765/mcp", type="http") - - with patch("cursor_sdk.AsyncClient") as mock_client_cls, patch.object( - backend, "_ensure_mcp_server", new_callable=AsyncMock - ), patch.object( - HttpMcpServer, "config", new_callable=PropertyMock, return_value=mock_mcp - ): - mock_client_cls.launch_bridge = AsyncMock(return_value=mock_client) - - result = await backend.run_turn( - system_prompt="System", - context="ctx", - user_message="User", - tool_executor=MagicMock(), - session_key="cli", - max_turns=5, - ) + with _patched(backend, client): + result = await _turn(backend) assert result.text == "Turn response" - assert result.num_turns == 2 + assert result.session_id == "agent-1" assert result.cost_usd is None - send_options = mock_agent.send.await_args.args[1] + create_options = client.agents.create.await_args.args[0] + assert create_options.mcp_servers == {"memclaw": _MCP} + assert create_options.model.id == "composer-2.5" + _, send_options = agent.sent[0] assert send_options.on_delta is not None - mock_client.agents.create.assert_awaited_once() - create_options = mock_client.agents.create.await_args.args[0] - assert create_options.mcp_servers == {"memclaw": mock_mcp} - send_options = mock_agent.send.await_args.args[1] - assert send_options.mcp_servers == {"memclaw": mock_mcp} - mock_agent.send.assert_awaited_once() - mock_agent.close.assert_awaited_once() + assert send_options.mcp_servers == {"memclaw": _MCP} + assert send_options.model.id == "composer-2.5" + assert send_options.local is None + agent.close.assert_not_awaited() # kept open for the next turn @pytest.mark.asyncio - async def test_run_turn_applies_cost_fallback(self, tmp_path): - cfg = _make_config(tmp_path, cursor_api_key="crsr_test_key") - backend = CursorAgentBackend(cfg) + async def test_agent_is_reused_and_system_prompt_only_sent_first(self, tmp_path): + backend = _backend(tmp_path) + agent = _make_agent(replies=["First", "Second"]) + client = _make_client(create=[agent]) + + with _patched(backend, client) as launch: + first = await _turn(backend, "My name is Ana") + second = await _turn(backend, "What's my name?", resume=first.session_id) + + launch.assert_awaited_once() # one bridge for the process + client.agents.create.assert_awaited_once() + client.agents.resume.assert_not_awaited() + assert second.session_id == "agent-1" + first_msg, second_msg = agent.sent[0][0], agent.sent[1][0] + assert first_msg.startswith("System") + assert "ctx" in first_msg and "My name is Ana" in first_msg + assert "System" not in second_msg + assert second_msg == "ctx\n\nWhat's my name?" - async def _events(): - yield SimpleNamespace(sdk_message=None, interaction_update=None, result=None) + @pytest.mark.asyncio + async def test_sessions_are_separate_per_key(self, tmp_path): + backend = _backend(tmp_path) + a, b = _make_agent("agent-a"), _make_agent("agent-b") + client = _make_client(create=[a, b]) - mock_run = AsyncMock() - mock_run.events = MagicMock(return_value=_events()) - mock_run.wait = AsyncMock(return_value=SimpleNamespace(result="Turn response", num_turns=1)) + with _patched(backend, client): + ra = await _turn(backend, key="chat-a") + rb = await _turn(backend, key="chat-b") + + assert (ra.session_id, rb.session_id) == ("agent-a", "agent-b") + + @pytest.mark.asyncio + async def test_resume_passes_model_and_forces_first_send(self, tmp_path): + backend = _backend(tmp_path, cursor_model="gpt-6") + agent = _make_agent("agent-old", replies=["One", "Two"]) + client = _make_client(resume=[agent]) + + with _patched(backend, client): + result = await _turn(backend, resume="agent-old") + await _turn(backend, resume="agent-old") + + client.agents.create.assert_not_awaited() + agent_id, options = client.agents.resume.await_args.args + assert agent_id == "agent-old" + assert options.model.id == "gpt-6" + assert options.mcp_servers == {"memclaw": _MCP} + assert options.local.cwd == backend._cwd + assert result.session_id == "agent-old" + # No system prompt for a resumed agent; it already has one. + assert "System" not in agent.sent[0][0] + assert agent.sent[0][1].local.force is True + assert agent.sent[1][1].local is None + + @pytest.mark.asyncio + async def test_resume_failure_starts_fresh_agent(self, tmp_path): + from cursor_sdk import UnknownAgentError + + backend = _backend(tmp_path) + fresh = _make_agent("agent-new") + client = _make_client( + resume=[UnknownAgentError("Agent agent-gone not found")], create=[fresh], + ) + + with _patched(backend, client): + result = await _turn(backend, resume="agent-gone") + + assert result.session_id == "agent-new" + assert fresh.sent[0][0].startswith("System") + + @pytest.mark.asyncio + async def test_failed_turn_drops_agent_and_resumes_next_time(self, tmp_path): + from cursor_sdk import CursorSDKError + + backend = _backend(tmp_path) + broken = _make_agent("agent-1") + broken.send = AsyncMock(side_effect=CursorSDKError("boom")) + again = _make_agent("agent-1") + client = _make_client(create=[broken], resume=[again]) + + with _patched(backend, client) as launch: + with pytest.raises(RuntimeError, match="Cursor SDK error: boom"): + await _turn(backend) + broken.close.assert_awaited_once() + result = await _turn(backend, resume="agent-1") + + assert result.session_id == "agent-1" + launch.assert_awaited_once() # the bridge itself was fine + client.aclose.assert_not_awaited() + + @pytest.mark.asyncio + async def test_dead_bridge_is_relaunched(self, tmp_path): + import httpx + from cursor_sdk import NetworkError + + backend = _backend(tmp_path) + dead = _make_agent("agent-1") + err = NetworkError("connection refused") + err.__cause__ = httpx.ConnectError("refused") + dead.send = AsyncMock(side_effect=err) + old_client = _make_client(create=[dead]) + new_client = _make_client(resume=[_make_agent("agent-1")]) + + with _patched(backend, old_client, new_client) as launch: + with pytest.raises(RuntimeError): + await _turn(backend) + await _turn(backend, resume="agent-1") + + assert launch.await_count == 2 + old_client.aclose.assert_awaited_once() + new_client.agents.resume.assert_awaited_once() + + @pytest.mark.asyncio + async def test_reset_session_closes_agent(self, tmp_path): + backend = _backend(tmp_path) + first, second = _make_agent("agent-1"), _make_agent("agent-2") + client = _make_client(create=[first, second]) + + with _patched(backend, client): + await _turn(backend) + await backend.reset_session("cli") + result = await _turn(backend) # agent.py passes no resume id after /new + + first.close.assert_awaited_once() + assert result.session_id == "agent-2" + assert second.sent[0][0].startswith("System") + + @pytest.mark.asyncio + async def test_reset_unknown_session_is_noop(self, tmp_path): + await _backend(tmp_path).reset_session("nope") + + @pytest.mark.asyncio + async def test_shutdown_closes_agents_and_bridge(self, tmp_path): + backend = _backend(tmp_path) + agent = _make_agent() + client = _make_client(create=[agent]) + + with _patched(backend, client), patch.object( + backend._mcp_server, "stop", new_callable=AsyncMock + ) as stop_mcp: + await _turn(backend) + await backend.on_agent_shutdown() + + agent.close.assert_awaited_once() + client.aclose.assert_awaited_once() + stop_mcp.assert_awaited_once() + assert backend._client is None + + @pytest.mark.asyncio + async def test_run_turn_applies_cost_fallback(self, tmp_path): + backend = _backend(tmp_path) + agent = _make_agent() async def _send(message, options): options.on_delta( @@ -398,30 +611,13 @@ async def _send(message, options): }, ) ) - return mock_run - - mock_agent = AsyncMock() - mock_agent.send = AsyncMock(side_effect=_send) - mock_agent.close = AsyncMock() - mock_client = AsyncMock() - mock_client.agents.create = AsyncMock(return_value=mock_agent) - mock_mcp = HttpMcpServerConfig(url="http://127.0.0.1:8765/mcp", type="http") - - with patch("cursor_sdk.AsyncClient") as mock_client_cls, patch.object( - backend, "_ensure_mcp_server", new_callable=AsyncMock - ), patch.object( - HttpMcpServer, "config", new_callable=PropertyMock, return_value=mock_mcp - ): - mock_client_cls.launch_bridge = AsyncMock(return_value=mock_client) - - result = await backend.run_turn( - system_prompt="System", - context="ctx", - user_message="User", - tool_executor=MagicMock(), - session_key="cli", - max_turns=5, - ) + return _make_run() + + agent.send = AsyncMock(side_effect=_send) + client = _make_client(create=[agent]) + + with _patched(backend, client): + result = await _turn(backend) assert result.input_tokens == 1000 assert result.output_tokens == 500 @@ -432,76 +628,13 @@ async def _send(message, options): @pytest.mark.asyncio async def test_run_turn_without_mcp_server_raises(self, tmp_path): - cfg = _make_config(tmp_path, cursor_api_key="crsr_test_key") - backend = CursorAgentBackend(cfg) + backend = _backend(tmp_path) with patch.object(backend, "_ensure_mcp_server", new_callable=AsyncMock), patch.object( HttpMcpServer, "config", new_callable=PropertyMock, return_value=None ): with pytest.raises(RuntimeError, match="MCP server failed to start"): - await backend.run_turn( - system_prompt="System", - context="", - user_message="User", - tool_executor=MagicMock(), - session_key="cli", - ) - - @pytest.mark.asyncio - async def test_interim_transcript_is_folded_into_later_prompts(self, tmp_path): - cfg = _make_config(tmp_path, cursor_api_key="crsr_test_key") - backend = CursorAgentBackend(cfg) - - sent: list[str] = [] - - def _make_run(reply): - run = AsyncMock() - run.events = MagicMock(return_value=_empty_events()) - run.wait = AsyncMock(return_value=SimpleNamespace(result=reply, num_turns=1)) - return run - - replies = iter(["First reply", "Second reply", "Third reply"]) - - async def _send(message, options): - sent.append(message) - return _make_run(next(replies)) - - mock_agent = AsyncMock() - mock_agent.send = AsyncMock(side_effect=_send) - mock_agent.close = AsyncMock() - mock_client = AsyncMock() - mock_client.agents.create = AsyncMock(return_value=mock_agent) - mock_mcp = HttpMcpServerConfig(url="http://127.0.0.1:8765/mcp", type="http") - - async def _turn(key, text): - return await backend.run_turn( - system_prompt="System", context="ctx", - user_message=text, tool_executor=MagicMock(), session_key=key, - ) - - with patch("cursor_sdk.AsyncClient") as mock_client_cls, patch.object( - backend, "_ensure_mcp_server", new_callable=AsyncMock - ), patch.object( - HttpMcpServer, "config", new_callable=PropertyMock, return_value=mock_mcp - ): - mock_client_cls.launch_bridge = AsyncMock(return_value=mock_client) - await _turn("cli", "My name is Ana") - await _turn("cli", "What's my name?") - await backend.reset_session("cli") - await _turn("cli", "Fresh start") - - assert "RECENT CONVERSATION" not in sent[0] - assert "ctx" in sent[0] - assert "User: My name is Ana" in sent[1] - assert "Assistant: First reply" in sent[1] - # The context block itself isn't kept in the transcript. - assert sent[1].count("ctx") == 1 - assert "RECENT CONVERSATION" not in sent[2] - - def test_model_selection_is_not_supported_yet(self, tmp_path): - backend = CursorAgentBackend(_make_config(tmp_path, cursor_api_key="k")) - assert backend.supports_model_selection is False - assert backend.model == "composer-2.5" + await _turn(backend) @pytest.mark.asyncio async def test_agent_start_starts_mcp_server(self, tmp_path, monkeypatch): @@ -527,6 +660,134 @@ async def test_agent_start_starts_mcp_server(self, tmp_path, monkeypatch): mock_stop.assert_awaited_once() +class TestCursorModelSelection: + @pytest.fixture(autouse=True) + def _env_file(self, tmp_path, monkeypatch): + self.env_file = tmp_path / ".env" + monkeypatch.setattr("memclaw.setup.ENV_FILE", self.env_file) + + def _env(self) -> str: + return self.env_file.read_text() if self.env_file.exists() else "" + + def test_supports_model_selection(self, tmp_path): + backend = _backend(tmp_path) + assert backend.supports_model_selection is True + assert backend.model == "composer-2.5" + assert backend.effort is None + + @pytest.mark.asyncio + async def test_list_models_maps_sdk_models(self, tmp_path): + backend = _backend(tmp_path) + client = _make_client() + client.models.list.return_value = [ + _sdk_model("composer-2.5"), + _sdk_model("gpt-6", effort_values=("low", "medium", "high")), + ] + + with _patched(backend, client): + models = await backend.list_models() + + client.models.list.assert_awaited_once_with(api_key="crsr_test_key") + assert [(m.id, m.display_name, m.effort_levels) for m in models] == [ + ("composer-2.5", "Composer-2.5", []), + ("gpt-6", "Gpt-6", ["low", "medium", "high"]), + ] + + @pytest.mark.asyncio + async def test_set_model_persists_and_applies_on_next_send(self, tmp_path): + backend = _backend(tmp_path) + agent = _make_agent(replies=["One", "Two"]) + client = _make_client(create=[agent]) + client.models.list.return_value = [_sdk_model("composer-2.5"), _sdk_model("gpt-6")] + + with _patched(backend, client): + await _turn(backend) + assert await backend.set_model("gpt-6") == "" + await _turn(backend) + + assert backend.model == "gpt-6" + assert "CURSOR_MODEL=gpt-6" in self._env() + assert agent.sent[0][1].model.id == "composer-2.5" + assert agent.sent[1][1].model.id == "gpt-6" + client.agents.create.assert_awaited_once() # same agent, no reconnect + + @pytest.mark.asyncio + async def test_set_unknown_model_raises(self, tmp_path): + backend = _backend(tmp_path) + client = _make_client() + client.models.list.return_value = [_sdk_model("composer-2.5")] + + with _patched(backend, client), pytest.raises(ValueError, match="Unknown model"): + await backend.set_model("nope") + + @pytest.mark.asyncio + async def test_effort_is_sent_as_model_parameter(self, tmp_path): + backend = _backend(tmp_path) + agent = _make_agent() + client = _make_client(create=[agent]) + client.models.list.return_value = [ + _sdk_model("gpt-6", effort_values=("low", "high")), + ] + + with _patched(backend, client): + await backend.set_model("gpt-6") + assert await backend.effort_levels() == ["low", "high"] + await backend.set_effort("HIGH") + await _turn(backend) + + assert backend.effort == "high" + assert "CURSOR_EFFORT=high" in self._env() + selection = agent.sent[0][1].model + assert selection.id == "gpt-6" + assert [(p.id, p.value) for p in selection.params] == [("reasoning", "high")] + + @pytest.mark.asyncio + async def test_effort_rejected_when_model_has_none(self, tmp_path): + backend = _backend(tmp_path) + client = _make_client() + client.models.list.return_value = [_sdk_model("composer-2.5")] + + with _patched(backend, client): + assert await backend.effort_levels() == [] + with pytest.raises(ValueError, match="doesn't support an effort"): + await backend.set_effort("high") + + @pytest.mark.asyncio + async def test_switching_to_model_without_effort_clears_it(self, tmp_path): + backend = _backend(tmp_path, cursor_effort="high") + client = _make_client() + client.models.list.return_value = [ + _sdk_model("gpt-6", effort_values=("low", "high")), + _sdk_model("composer-2.5"), + ] + + with _patched(backend, client): + note = await backend.set_model("composer-2.5") + + assert "doesn't support 'high'" in note + assert backend.effort is None + assert "CURSOR_EFFORT" not in self._env() + + @pytest.mark.asyncio + async def test_model_and_effort_menus_list_cursor_models(self, tmp_path): + from memclaw.bot.handlers import effort_menu, model_menu, select_model + + backend = _backend(tmp_path) + client = _make_client() + client.models.list.return_value = [_sdk_model("composer-2.5"), _sdk_model("gpt-6")] + + with _patched(backend, client): + _text, keyboard = await model_menu(backend) + switched = await select_model(backend, "gpt-6") + effort, effort_keyboard = await effort_menu(backend) + + buttons = [(b.text, b.callback_data) for row in keyboard.inline_keyboard for b in row] + assert buttons == [("✓ Composer-2.5", "m:composer-2.5"), ("Gpt-6", "m:gpt-6")] + assert switched.startswith("Switched to gpt-6.") + assert "doesn't support an effort setting" in effort + assert effort_keyboard is None + + async def _empty_events(): if False: # pragma: no cover - async generator helper yield diff --git a/tests/test_cli_status.py b/tests/test_cli_status.py index ad54208..5936ab4 100644 --- a/tests/test_cli_status.py +++ b/tests/test_cli_status.py @@ -37,13 +37,15 @@ def test_claude_backend_shows_its_model_and_effort(tmp_path: Path, monkeypatch): def test_cursor_backend_shows_only_its_own_model(tmp_path: Path, monkeypatch): monkeypatch.setenv("AGENT_BACKEND", "cursor") monkeypatch.setenv("CURSOR_MODEL", "composer-3") + monkeypatch.delenv("CURSOR_EFFORT", raising=False) monkeypatch.setenv("CLAUDE_MODEL", "claude-opus-5") + monkeypatch.setenv("CLAUDE_EFFORT", "high") output = _status(tmp_path) assert "Model : composer-3" in output + assert "Effort : default" in output assert "claude-opus-5" not in output - assert "Effort" not in output def test_general_rows_are_still_shown(tmp_path: Path): From a6f1b1333e270c601aa14ac094deda173d5ec32e Mon Sep 17 00:00:00 2001 From: Emanuilo Jovanovic Date: Sun, 18 Oct 2026 11:56:43 +0200 Subject: [PATCH 05/10] fix: review findings - /new forgets the session before awaiting the backend reset, so a message sent right after it can't resume the old session - /model refuses to switch when the model list can't be loaded, instead of saving an unchecked (possibly mistyped) id that breaks every turn - the fallback cost estimate includes cache-write tokens - long-id model buttons carry a hash of the id instead of a list position, so they can't pick a different model if the list changed --- memclaw/agent.py | 5 ++++- memclaw/backends/claude.py | 10 ++++++++-- memclaw/bot/handlers.py | 23 +++++++++++++++-------- tests/backends/test_claude.py | 11 +++++++---- tests/test_agent.py | 23 +++++++++++++++++++++++ tests/test_handlers.py | 16 ++++++++++------ 6 files changed, 67 insertions(+), 21 deletions(-) diff --git a/memclaw/agent.py b/memclaw/agent.py index 3792b59..19d4f79 100644 --- a/memclaw/agent.py +++ b/memclaw/agent.py @@ -210,9 +210,12 @@ async def reset_conversation(self, chat_id: str | None = None) -> None: """Start a fresh conversation for *chat_id* (the `/new` command).""" key = self._session_key(chat_id) self._generations[key] = self._generations.get(key, 0) + 1 - await self.backend.reset_session(key) + # Forget the session before awaiting the backend: a message sent + # right after /new may run while the reset is still disconnecting, + # and must not resume the old session. self.sessions.drop(key) self._pending_notes.pop(key, None) + await self.backend.reset_session(key) logger.info("Started a new conversation for {k}", k=key) # ── Startup / sync ─────────────────────────────────────────────── diff --git a/memclaw/backends/claude.py b/memclaw/backends/claude.py index 0a65cd0..e2fe90b 100644 --- a/memclaw/backends/claude.py +++ b/memclaw/backends/claude.py @@ -68,8 +68,10 @@ "claude-haiku-4-5": (1.0, 5.0), } -# Cache reads are billed at a tenth of the input price. +# Cache reads are billed at a tenth of the input price, cache writes at +# 1.25 times it. _CACHE_READ_MULTIPLIER = 0.1 +_CACHE_WRITE_MULTIPLIER = 1.25 _ALLOWED_TOOLS = [ f"mcp__{MCP_SERVER_NAME}__{t['name']}" for t in TOOL_DEFINITIONS @@ -580,7 +582,10 @@ async def effort_levels(self) -> list[str] | None: async def set_model(self, model_id: str) -> str: model_id = model_id.strip() info = await self._model_info(model_id) - if info is None and self._models: + if info is None and not self._models: + # Can't check the id, and a typo would break every turn. + raise ValueError("Couldn't load the model list; try again later.") + if info is None: raise ValueError(f"Unknown model '{model_id}'.") effort = self._effort note = "" @@ -724,6 +729,7 @@ async def run_turn( result.input_tokens * input_per_m + result.output_tokens * output_per_m + result.cache_read_tokens * input_per_m * _CACHE_READ_MULTIPLIER + + result.cache_creation_tokens * input_per_m * _CACHE_WRITE_MULTIPLIER ) / 1_000_000 return result diff --git a/memclaw/bot/handlers.py b/memclaw/bot/handlers.py index 46e0623..e67a4c0 100644 --- a/memclaw/bot/handlers.py +++ b/memclaw/bot/handlers.py @@ -9,6 +9,7 @@ import asyncio import base64 import contextlib +import hashlib import time from loguru import logger @@ -57,10 +58,16 @@ async def _typing_loop(bot, chat_id: int): _APPLIES = "Applies from your next message; the conversation is kept." -def model_callback_data(index: int, model_id: str) -> str: - """``m:``, or ``mi:`` if the id would exceed Telegram's limit.""" +def _model_key(model_id: str) -> str: + return hashlib.sha256(model_id.encode()).hexdigest()[:16] + + +def model_callback_data(model_id: str) -> str: + """``m:``, or ``mi:`` if the id would exceed + Telegram's limit. A hash rather than a list position, so a button still + picks its own model if the list has changed since the menu was shown.""" data = f"m:{model_id}" - return data if len(data.encode()) <= CALLBACK_DATA_MAX_BYTES else f"mi:{index}" + return data if len(data.encode()) <= CALLBACK_DATA_MAX_BYTES else f"mi:{_model_key(model_id)}" def _unsupported(backend: AgentBackend, what: str) -> str: @@ -97,15 +104,15 @@ async def model_menu( lines += ["", note] lines.append("") if not models: - lines.append("Couldn't load the model list right now; send /model to switch.") + lines.append("Couldn't load the model list right now; try /model again later.") return "\n".join(lines), None lines.append("Tap a model to switch, or send /model .") keyboard = InlineKeyboardMarkup([ [InlineKeyboardButton( ("✓ " if m.id == backend.model else "") + m.display_name, - callback_data=model_callback_data(i, m.id), + callback_data=model_callback_data(m.id), )] - for i, m in enumerate(models) + for m in models ]) return "\n".join(lines), keyboard @@ -338,9 +345,9 @@ async def handle_callback(self, update: Update, context: ContextTypes.DEFAULT_TY model_id = value if kind == "mi": models = await _load_models(backend) - if not (value.isdigit() and int(value) < len(models)): + model_id = next((m.id for m in models if _model_key(m.id) == value), "") + if not model_id: raise ValueError("That model is no longer available; send /model again.") - model_id = models[int(value)].id note = await select_model(backend, model_id) await query.answer(f"Model: {backend.model}") text, keyboard = await model_menu(backend, note) diff --git a/tests/backends/test_claude.py b/tests/backends/test_claude.py index b8fad97..4123dbb 100644 --- a/tests/backends/test_claude.py +++ b/tests/backends/test_claude.py @@ -506,15 +506,18 @@ async def test_unknown_model_is_rejected(self, claude_config, env_file): assert backend.model == "claude-opus-5" @pytest.mark.asyncio - async def test_any_id_is_accepted_when_the_list_is_unavailable( + async def test_switch_is_refused_when_the_list_is_unavailable( self, claude_config, env_file, ): + """An unchecked id (e.g. a typo) would break every turn until .env + is fixed by hand, so nothing is switched or saved.""" backend = self._backend(claude_config) + before = backend.model, backend.effort with patch("memclaw.backends.claude.fetch_models", AsyncMock(side_effect=httpx.ConnectError("offline"))): - await backend.set_model("claude-future-9") - assert backend.model == "claude-future-9" - assert backend.effort == "high" + with pytest.raises(ValueError, match="Couldn't load the model list"): + await backend.set_model("claude-opsu-5") + assert (backend.model, backend.effort) == before @pytest.mark.asyncio async def test_set_effort(self, claude_config, env_file): diff --git a/tests/test_agent.py b/tests/test_agent.py index 2b50443..6e33f8f 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -724,3 +724,26 @@ def test_tool_definitions_contain_all_tools(self): } assert set(names) == expected assert len(names) == len(expected) + + + +class TestNewDuringReset: + @pytest.mark.asyncio + async def test_message_right_after_new_does_not_resume_the_old_session( + self, cfg: MemclawConfig, fake_backend: FakeBackend, + ): + """A message that runs while /new is still disconnecting the backend + must start fresh, not resume (and re-save) the old session.""" + agent = _chat_agent(cfg, fake_backend, platform="slack") + await agent.handle("hi", chat_id="C1") + seen: list[str | None] = [] + + async def slow_reset(key): + # The next message arrives while the reset is in progress. + await agent.handle("next", chat_id="C1") + seen.append(fake_backend.turn_calls[-1]["resume_session_id"]) + + fake_backend.reset_session = slow_reset + await agent.reset_conversation("C1") + assert seen == [None] + agent.close() diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 5bf7401..33e2926 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -3,7 +3,6 @@ import asyncio import inspect -import textwrap from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -341,11 +340,12 @@ async def test_menu_has_a_button_per_model(self): ("Claude Haiku 4.5", "m:claude-haiku-4-5"), ] - def test_long_model_id_falls_back_to_index(self): - from memclaw.bot.handlers import model_callback_data + def test_long_model_id_falls_back_to_a_hash(self): + from memclaw.bot.handlers import _model_key, model_callback_data - assert model_callback_data(3, "x" * 70) == "mi:3" - assert model_callback_data(3, "claude-opus-5") == "m:claude-opus-5" + assert model_callback_data("x" * 70) == f"mi:{_model_key('x' * 70)}" + assert len(model_callback_data("x" * 70).encode()) <= 64 + assert model_callback_data("claude-opus-5") == "m:claude-opus-5" @pytest.mark.asyncio async def test_button_press_switches_and_edits_the_menu(self): @@ -368,7 +368,11 @@ async def test_button_press_switches_and_edits_the_menu(self): async def test_index_button_press(self): handlers = _telegram_handlers() handlers.agent.backend = backend = _SelectableBackend() - await handlers.handle_callback(_callback_update("mi:1"), MagicMock()) + from memclaw.bot.handlers import _model_key + + await handlers.handle_callback( + _callback_update(f"mi:{_model_key('claude-haiku-4-5')}"), MagicMock(), + ) assert backend.model == "claude-haiku-4-5" @pytest.mark.asyncio From ac8c1a4209fc82f12f7f47677f80f64694b31187 Mon Sep 17 00:00:00 2001 From: Emanuilo Jovanovic Date: Sun, 18 Oct 2026 11:58:40 +0200 Subject: [PATCH 06/10] fix: Cursor bridge robustness (review findings) - drain the bridge's stderr in chunks: readline() raised on a line over 64 KiB, ending the drain, after which a full pipe would block the bridge - a timeout no longer counts as a dead bridge, so one slow run doesn't drop the shared bridge and abort other chats' runs - relaunch a bridge that exited between calls instead of failing the next message (or /model) first - shutdown iterates over a copy of the sessions, so a message arriving mid-shutdown can't skip closing the bridge and MCP server --- memclaw/backends/cursor.py | 33 ++++++++++++++---- tests/backends/test_cursor.py | 64 ++++++++++++++++++++++++++++++++++- 2 files changed, 89 insertions(+), 8 deletions(-) diff --git a/memclaw/backends/cursor.py b/memclaw/backends/cursor.py index 9c2da99..e7804b7 100644 --- a/memclaw/backends/cursor.py +++ b/memclaw/backends/cursor.py @@ -148,15 +148,25 @@ def _effort_param(model: Any) -> Any | None: async def _drain_stderr(stream: asyncio.StreamReader) -> None: """Keep reading the bridge's stderr. The SDK only reads it until the - bridge is up, and a long-lived bridge would block once the pipe fills.""" - while line := await stream.readline(): - logger.debug("cursor-sdk-bridge: {line}", line=line.decode(errors="replace").rstrip()) + bridge is up, and a long-lived bridge would block once the pipe fills. + Read in chunks: readline() fails on a line over 64 KiB, which would end + the drain.""" + while chunk := await stream.read(65536): + logger.debug("cursor-sdk-bridge: {out}", out=chunk.decode(errors="replace").rstrip()) def _bridge_failed(exc: BaseException) -> bool: """True when *exc* came from not reaching the bridge process at all - (it crashed or hung), rather than from the agent or Cursor's API.""" - return isinstance(exc.__cause__, httpx.RequestError) + (it crashed), rather than from the agent or Cursor's API. A timeout + doesn't count: a slow run says nothing about the bridge, and dropping it + would abort other chats' runs.""" + cause = exc.__cause__ + return isinstance(cause, httpx.RequestError) and not isinstance(cause, httpx.TimeoutException) + + +def _bridge_process(client: Any) -> asyncio.subprocess.Process | None: + # The SDK keeps the subprocess it launched on a private attribute. + return getattr(getattr(client, "_owned_bridge", None), "process", None) @dataclass @@ -282,7 +292,8 @@ async def on_agent_start(self, tool_executor: "ToolExecutor") -> None: ) async def on_agent_shutdown(self) -> None: - for session in self._sessions.values(): + # A copy: a message arriving mid-shutdown may add a session. + for session in list(self._sessions.values()): await self._close_agent(session) self._sessions.clear() if self._client is not None: @@ -310,7 +321,7 @@ async def _launch_client(self) -> Any: ["/bin/sh", "-c", 'cd "$0" && exec "$@"', self._cwd, resolve_bridge_path()], workspace=self._cwd, ) - process = getattr(client._owned_bridge, "process", None) + process = _bridge_process(client) if process is not None and process.stderr is not None: task = asyncio.create_task(_drain_stderr(process.stderr)) self._background.add(task) @@ -319,6 +330,14 @@ async def _launch_client(self) -> Any: async def _get_client(self) -> Any: async with self._client_lock: + client = self._client + if client is not None: + process = _bridge_process(client) + if process is not None and process.returncode is not None: + # The bridge died between calls: relaunch now rather than + # failing the next message (or /model) first. + logger.warning("Cursor SDK bridge exited ({rc}); relaunching", rc=process.returncode) + await self._drop_client(client) if self._client is None: self._client = await self._launch_client() return self._client diff --git a/tests/backends/test_cursor.py b/tests/backends/test_cursor.py index a21f547..b738d6e 100644 --- a/tests/backends/test_cursor.py +++ b/tests/backends/test_cursor.py @@ -2,12 +2,14 @@ from __future__ import annotations +import asyncio import os from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch +import httpx import pytest from cursor_sdk import ( HttpMcpServerConfig, @@ -342,8 +344,9 @@ async def _send(message, options): return agent -def _make_client(*, create=None, resume=None): +def _make_client(*, create=None, resume=None, process=None): client = MagicMock() + client._owned_bridge = MagicMock(process=process) if process else None client.agents.create = AsyncMock(side_effect=create or [_make_agent()]) client.agents.resume = AsyncMock(side_effect=resume or []) client.models.list = AsyncMock(return_value=[]) @@ -660,6 +663,65 @@ async def test_agent_start_starts_mcp_server(self, tmp_path, monkeypatch): mock_stop.assert_awaited_once() +class TestBridgeRobustness: + @pytest.mark.asyncio + async def test_dead_bridge_is_relaunched_before_the_next_call(self, tmp_path): + backend = _backend(tmp_path) + dead = _make_client(process=MagicMock(returncode=1)) + fresh = _make_client() + with patch.object( + backend, "_launch_client", new_callable=AsyncMock, side_effect=[dead, fresh], + ): + assert await backend._get_client() is dead + assert await backend._get_client() is fresh + dead.aclose.assert_awaited_once() + + @pytest.mark.asyncio + async def test_running_bridge_is_reused(self, tmp_path): + backend = _backend(tmp_path) + alive = _make_client(process=MagicMock(returncode=None)) + with patch.object(backend, "_launch_client", new_callable=AsyncMock, side_effect=[alive]): + assert await backend._get_client() is alive + assert await backend._get_client() is alive + + def test_timeout_is_not_a_dead_bridge(self): + from memclaw.backends.cursor import _bridge_failed + + def raised_from(cause): + exc = RuntimeError("run failed") + exc.__cause__ = cause + return exc + + assert _bridge_failed(raised_from(httpx.ConnectError("refused"))) + assert not _bridge_failed(raised_from(httpx.ReadTimeout("slow"))) + assert not _bridge_failed(RuntimeError("agent error")) + + @pytest.mark.asyncio + async def test_drain_survives_a_line_over_64_kib(self): + from memclaw.backends.cursor import _drain_stderr + + stream = asyncio.StreamReader() + stream.feed_data(b"x" * 200_000 + b"\nok\n") + stream.feed_eof() + await _drain_stderr(stream) # reads to EOF without raising + assert stream.at_eof() + + @pytest.mark.asyncio + async def test_shutdown_tolerates_a_session_added_meanwhile(self, tmp_path): + from memclaw.backends.cursor import _Session + + backend = _backend(tmp_path) + backend._mcp_server = MagicMock(stop=AsyncMock()) + + async def close_and_add(session): + backend._sessions.setdefault("late", _Session()) + + backend._sessions["a"] = _Session() + with patch.object(backend, "_close_agent", side_effect=close_and_add): + await backend.on_agent_shutdown() + assert backend._sessions == {} + + class TestCursorModelSelection: @pytest.fixture(autouse=True) def _env_file(self, tmp_path, monkeypatch): From 86fafa4b05dce4740b064053bba570a896b5b2b6 Mon Sep 17 00:00:00 2001 From: Emanuilo Jovanovic Date: Sun, 18 Oct 2026 18:50:39 +0200 Subject: [PATCH 07/10] refactor: tidy the Claude backend and agent.handle - one builder for the locked-down CLI options, shared by chat sessions and one-shots, so the lockdown can't drift between them - turn cost computed in _turn_cost, out of run_turn's connection handling - agent.handle's logging, pending-note acknowledgement and session saving pulled into named helpers - clear the repo's ruff findings (unused imports, ambiguous names) --- memclaw/agent.py | 62 +++++++++++++++++++-------------- memclaw/backends/claude.py | 71 +++++++++++++++++++------------------- memclaw/config.py | 1 - memclaw/index.py | 4 +-- tests/conftest.py | 3 -- tests/test_index.py | 2 +- tests/test_search.py | 2 -- 7 files changed, 75 insertions(+), 70 deletions(-) diff --git a/memclaw/agent.py b/memclaw/agent.py index 19d4f79..e8732bc 100644 --- a/memclaw/agent.py +++ b/memclaw/agent.py @@ -24,7 +24,7 @@ from loguru import logger from .backends import AgentBackend, build_backend -from .backends.base import ProgressCallback +from .backends.base import ProgressCallback, TurnResult from .config import MemclawConfig from .index import MemoryIndex from .reminders import ReminderScheduler @@ -456,48 +456,58 @@ async def handle( max_turns=10, on_tool=on_tool, ) - elapsed_ms = int((time.perf_counter() - t0) * 1000) + self._log_turn(result, int((time.perf_counter() - t0) * 1000)) - token_summary = ( + # The session has now seen these notes and these file versions. Its + # own writes during the turn (memory_save permanent=true, + # update_instructions) count as seen too, so they aren't sent back + # next turn; a consolidation rewrite during the turn doesn't. + agents_hash = _hash(_load_agent_instructions(self.config)) + if self._memory_rewrites == memory_rewrites: + memory_hash = _hash(_load_permanent_memory(self.config)) + self._ack_notes(key, len(notes)) + if result.session_id and self._generations.get(key, 0) == generation: + self._save_session(key, result.session_id, agents_hash, memory_hash) + + self._schedule_consolidation() + + response_text = result.text or "I couldn't generate a response." + return (response_text, list(self._found_images)) + + def _log_turn(self, result: TurnResult, elapsed_ms: int) -> None: + tokens = ( f"in={result.input_tokens}, out={result.output_tokens}, " f"cache_read={result.cache_read_tokens}, " f"cache_create={result.cache_creation_tokens}" ) + turns = result.num_turns or 1 if self.backend.bills_per_token and result.cost_usd is not None: logger.info( "Agent done: {turns} turns, {ms}ms, cost ${cost:.4f} ({tokens})", - turns=result.num_turns or 1, ms=elapsed_ms, - cost=result.cost_usd, tokens=token_summary, + turns=turns, ms=elapsed_ms, cost=result.cost_usd, tokens=tokens, ) else: # Subscription-billed backends, or backends that don't report cost. logger.info( "Agent done: {turns} turns, {ms}ms ({tokens})", - turns=result.num_turns or 1, ms=elapsed_ms, tokens=token_summary, + turns=turns, ms=elapsed_ms, tokens=tokens, ) - # The session has now seen these notes and these file versions. Its - # own writes during the turn (memory_save permanent=true, - # update_instructions) count as seen too, so they aren't sent back - # next turn; a consolidation rewrite during the turn doesn't. - agents_hash = _hash(_load_agent_instructions(self.config)) - if self._memory_rewrites == memory_rewrites: - memory_hash = _hash(_load_permanent_memory(self.config)) + def _ack_notes(self, key: str, seen: int) -> None: + """Drop the first *seen* pending notes, the ones this turn carried. + Only after the turn succeeded, and only those: a reminder delivered + during the turn stays queued for the next message.""" pending = self._pending_notes.get(key) if pending: - del pending[:len(notes)] - if result.session_id and self._generations.get(key, 0) == generation: - self.sessions.set(key, { - "session_id": result.session_id, - "fingerprint": self._fingerprint, - "agents_hash": agents_hash, - "memory_hash": memory_hash, - }) - - self._schedule_consolidation() - - response_text = result.text or "I couldn't generate a response." - return (response_text, list(self._found_images)) + del pending[:seen] + + def _save_session(self, key: str, session_id: str, agents_hash: str, memory_hash: str) -> None: + self.sessions.set(key, { + "session_id": session_id, + "fingerprint": self._fingerprint, + "agents_hash": agents_hash, + "memory_hash": memory_hash, + }) def _close_sync_only(self) -> None: for task in (getattr(self, "_sync_task", None), self._consolidation_task): diff --git a/memclaw/backends/claude.py b/memclaw/backends/claude.py index e2fe90b..b56bc6e 100644 --- a/memclaw/backends/claude.py +++ b/memclaw/backends/claude.py @@ -481,32 +481,41 @@ async def on_agent_shutdown(self) -> None: # -- Sessions ---------------------------------------------------------- - def _session_options( - self, *, system_prompt: str, resume: str | None, max_turns: int, - ) -> ClaudeAgentOptions: + def _options(self, *, system_prompt: str, max_turns: int, **extra: Any) -> ClaudeAgentOptions: + """The locked-down CLI options every call shares (chat sessions and + one-shots alike), plus *extra*.""" return ClaudeAgentOptions( env=self._env, model=self._model, effort=self._effort, system_prompt=system_prompt, - # Claude Code stores sessions per working directory; a fixed one - # lets a restarted process find (and resume) them. - cwd=str(self.config.memory_dir), # Don't load ~/.claude settings, plugins, CLAUDE.md or the user's # MCP servers: faster startup and nothing outside our control. setting_sources=[], strict_mcp_config=True, - mcp_servers={MCP_SERVER_NAME: self._mcp_server}, # No built-in tools at all; only our in-process mcp__memclaw__*. tools=[], - allowed_tools=_ALLOWED_TOOLS, # Anything not in allowed_tools is denied, never prompted for. permission_mode="dontAsk", # No @file expansion / slash commands from chat text. verbatim_prompts=True, max_turns=max_turns, - resume=resume, stderr=lambda line: logger.debug("claude-cli: {line}", line=line.rstrip()), + **extra, + ) + + def _session_options( + self, *, system_prompt: str, resume: str | None, max_turns: int, + ) -> ClaudeAgentOptions: + return self._options( + system_prompt=system_prompt, + max_turns=max_turns, + # Claude Code stores sessions per working directory; a fixed one + # lets a restarted process find (and resume) them. + cwd=str(self.config.memory_dir), + mcp_servers={MCP_SERVER_NAME: self._mcp_server}, + allowed_tools=_ALLOWED_TOOLS, + resume=resume, ) async def _connect( @@ -645,17 +654,7 @@ async def run_one_shot( system_prompt: str, user_message: str, ) -> str: - options = ClaudeAgentOptions( - env=self._env, - model=self._model, - effort=self._effort, - system_prompt=system_prompt, - setting_sources=[], - strict_mcp_config=True, - tools=[], - verbatim_prompts=True, - max_turns=1, - ) + options = self._options(system_prompt=system_prompt, max_turns=1) result_text = "" async with ClaudeSDKClient(options=options) as client: @@ -712,26 +711,28 @@ async def run_turn( await self._disconnect(session) raise - # Turn the session's running total into this turn's cost. Without a - # baseline (first turn after a restart resumed it) estimate instead. + result.cost_usd = self._turn_cost(result, resume_session_id) + return result + + def _turn_cost(self, result: TurnResult, resume_session_id: str | None) -> float | None: + """This turn's cost from the session's running total, which is what + the CLI reports. Without a baseline (the first turn after a restart + resumed the session) it is estimated from the token counts instead.""" total = result.cost_usd - result.cost_usd = None if total is not None and result.session_id: previous = self._session_costs.get(result.session_id) self._session_costs[result.session_id] = total if previous is not None or resume_session_id != result.session_id: - result.cost_usd = max(total - (previous or 0.0), 0.0) - - # Fallback cost estimate when the SDK didn't supply a usable one. - if result.cost_usd is None and self.bills_per_token: - input_per_m, output_per_m = _prices_per_m(self._model) - result.cost_usd = ( - result.input_tokens * input_per_m - + result.output_tokens * output_per_m - + result.cache_read_tokens * input_per_m * _CACHE_READ_MULTIPLIER - + result.cache_creation_tokens * input_per_m * _CACHE_WRITE_MULTIPLIER - ) / 1_000_000 - return result + return max(total - (previous or 0.0), 0.0) + if not self.bills_per_token: + return None + input_per_m, output_per_m = _prices_per_m(self._model) + return ( + result.input_tokens * input_per_m + + result.output_tokens * output_per_m + + result.cache_read_tokens * input_per_m * _CACHE_READ_MULTIPLIER + + result.cache_creation_tokens * input_per_m * _CACHE_WRITE_MULTIPLIER + ) / 1_000_000 @staticmethod async def _query( diff --git a/memclaw/config.py b/memclaw/config.py index 97ca69d..9dc087d 100644 --- a/memclaw/config.py +++ b/memclaw/config.py @@ -2,7 +2,6 @@ import importlib.resources import os -import shutil from dataclasses import dataclass, field from datetime import date from pathlib import Path diff --git a/memclaw/index.py b/memclaw/index.py index eb0edcc..a2d7297 100644 --- a/memclaw/index.py +++ b/memclaw/index.py @@ -191,10 +191,10 @@ def chunk_text(self, text: str, file_path: str) -> list[Chunk]: break current_lines = current_lines[overlap_start:] - if current_heading and not any(l.startswith("#") for l in current_lines): + if current_heading and not any(ln.startswith("#") for ln in current_lines): current_lines.insert(0, current_heading) current_start = max(0, i - len(current_lines) + 1) - word_count = sum(len(l.split()) for l in current_lines) + word_count = sum(len(ln.split()) for ln in current_lines) # Flush remaining content if current_lines: diff --git a/tests/conftest.py b/tests/conftest.py index da868f6..4662c04 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,10 +1,7 @@ """Shared fixtures for memclaw tests.""" from __future__ import annotations -import asyncio -from datetime import date from pathlib import Path -from unittest.mock import AsyncMock, MagicMock, patch import numpy as np import pytest diff --git a/tests/test_index.py b/tests/test_index.py index d95a06c..1d9d9b3 100644 --- a/tests/test_index.py +++ b/tests/test_index.py @@ -109,7 +109,7 @@ async def test_cache_miss_on_different_model(self, idx: MemoryIndex): mock_openai.embeddings.create = AsyncMock(return_value=mock_response) idx._openai = mock_openai - result = await idx.get_embeddings([text]) + await idx.get_embeddings([text]) # API should have been called because model doesn't match mock_openai.embeddings.create.assert_called_once() diff --git a/tests/test_search.py b/tests/test_search.py index 39723fd..7f06a2c 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -1,10 +1,8 @@ """Tests for HybridSearch — decay, MMR, file_filter, vector caching (specs #3–5, #8).""" from __future__ import annotations -import math from datetime import date, timedelta from pathlib import Path -from unittest.mock import patch import numpy as np import pytest From b1613668bd69aeb6fcb3033a50894c8684cd8594 Mon Sep 17 00:00:00 2001 From: Emanuilo Jovanovic Date: Sun, 18 Oct 2026 18:52:40 +0200 Subject: [PATCH 08/10] refactor: share model selection between backends; split Cursor run_turn - model/effort state, model-list lookups, the strict model check and saving to ~/.memclaw/.env move into ModelSelectionMixin, used by both the Claude and Cursor backends. Cursor's /model now also refuses an id it can't check when the model list doesn't load. - Cursor run_turn: _ensure_agent says whether the agent was already open, resumed or created (replacing two flags), and _send does the send --- memclaw/backends/claude.py | 80 ++++----------- memclaw/backends/cursor.py | 147 +++++++++++++--------------- memclaw/backends/model_selection.py | 84 ++++++++++++++++ tests/backends/test_cursor.py | 11 +++ 4 files changed, 180 insertions(+), 142 deletions(-) create mode 100644 memclaw/backends/model_selection.py diff --git a/memclaw/backends/claude.py b/memclaw/backends/claude.py index b56bc6e..470d48b 100644 --- a/memclaw/backends/claude.py +++ b/memclaw/backends/claude.py @@ -35,6 +35,7 @@ from .base import ProgressCallback, TurnResult, report_tool_step from .claude_models import SDK_EFFORT_LEVELS, ModelInfo, fetch_models from .mcp_tools import MCP_SERVER_NAME +from .model_selection import ModelSelectionMixin if TYPE_CHECKING: from rich.console import Console @@ -253,9 +254,11 @@ class _Session: options_key: tuple[str, str | None] | None = None -class ClaudeAgentBackend: +class ClaudeAgentBackend(ModelSelectionMixin): """Claude Agent SDK implementation of the AgentBackend protocol.""" + _label = "Claude" + _env_prefix = "CLAUDE" name: ClassVar[str] = "claude" display_name: ClassVar[str] = "Claude Agent SDK" supports_model_selection: ClassVar[bool] = True @@ -561,55 +564,26 @@ async def reset_session(self, session_key: str) -> None: # -- Model selection --------------------------------------------------- - @property - def model(self) -> str: - return self._model - - @property - def effort(self) -> str | None: - return self._effort - async def list_models(self) -> list[ModelInfo]: self._models = await fetch_models(_claude_auth_mode(self.config), _credential(self.config)) return self._models - async def _model_info(self, model_id: str) -> ModelInfo | None: - """*model_id*'s entry in the model list, or None if it isn't listed - or the list can't be fetched.""" - if not self._models: - try: - await self.list_models() - except Exception as exc: - logger.warning("Could not fetch the Claude model list: {exc}", exc=exc) - return None - return next((m for m in self._models if m.id == model_id), None) - - async def effort_levels(self) -> list[str] | None: - info = await self._model_info(self._model) - return list(info.effort_levels) if info else None - async def set_model(self, model_id: str) -> str: - model_id = model_id.strip() - info = await self._model_info(model_id) - if info is None and not self._models: - # Can't check the id, and a typo would break every turn. - raise ValueError("Couldn't load the model list; try again later.") - if info is None: - raise ValueError(f"Unknown model '{model_id}'.") - effort = self._effort + """Switch from the next turn on. Live chats then reconnect, resuming + their session (Claude Code resumes fine under a different model).""" + info = await self._known_model(model_id.strip()) + # Keep the effort if the new model takes it, otherwise the nearest + # level it does take (or none at all). + effort = _fit_effort(info.effort_levels, self._effort or _DEFAULT_EFFORT) note = "" - if info is not None: - # Keep the effort if the new model takes it, otherwise the nearest - # level it does take (or none at all). - effort = _fit_effort(info.effort_levels, self._effort or _DEFAULT_EFFORT) - if self._effort and effort is None: - note = f"{info.display_name} doesn't support an effort setting, so none is used." - elif self._effort and effort != self._effort: - note = ( - f"{info.display_name} doesn't support '{self._effort}' effort; " - f"using '{effort}' instead." - ) - self._apply_model(model_id, effort) + if self._effort and effort is None: + note = f"{info.display_name} doesn't support an effort setting, so none is used." + elif self._effort and effort != self._effort: + note = ( + f"{info.display_name} doesn't support '{self._effort}' effort; " + f"using '{effort}' instead." + ) + self._apply_model(info.id, effort) return note async def set_effort(self, level: str) -> None: @@ -628,24 +602,6 @@ async def set_effort(self, level: str) -> None: ) self._apply_model(self._model, level) - def _apply_model(self, model: str, effort: str | None) -> None: - """Use *model* / *effort* from the next turn on and save them to - ~/.memclaw/.env. Live chats reconnect on their next turn, resuming - their session (Claude Code resumes fine under a different model).""" - from ..setup import update_env_file # local import — setup imports backends - - logger.info( - "Claude model: {m} effort: {e} -> {m2} effort: {e2}", - m=self._model, e=self._effort, m2=model, e2=effort, - ) - self._model, self._effort = model, effort - self.config.claude_model = model - self.config.claude_effort = effort or "" - if effort: - update_env_file({"CLAUDE_MODEL": model, "CLAUDE_EFFORT": effort}) - else: - update_env_file({"CLAUDE_MODEL": model}, drop_keys=["CLAUDE_EFFORT"]) - # -- Runtime --------------------------------------------------------- async def run_one_shot( diff --git a/memclaw/backends/cursor.py b/memclaw/backends/cursor.py index e7804b7..a4b52fa 100644 --- a/memclaw/backends/cursor.py +++ b/memclaw/backends/cursor.py @@ -13,7 +13,7 @@ import os from dataclasses import dataclass, field from pathlib import Path -from typing import TYPE_CHECKING, Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar, Literal import httpx from loguru import logger @@ -23,6 +23,7 @@ from .cursor_hooks import cursor_hooks_status, ensure_cursor_hooks from .cursor_sdk_adapter import RunUsageTracker, collect_run_result, extract_run_text from .mcp_bridge import HttpMcpServer, mcp_servers_for +from .model_selection import ModelSelectionMixin if TYPE_CHECKING: from rich.console import Console @@ -180,9 +181,11 @@ class _Session: client: Any = None -class CursorAgentBackend: +class CursorAgentBackend(ModelSelectionMixin): """Cursor Python SDK implementation of the AgentBackend protocol.""" + _label = "Cursor" + _env_prefix = "CURSOR" name: ClassVar[str] = "cursor" display_name: ClassVar[str] = "Cursor SDK" supports_model_selection: ClassVar[bool] = True @@ -391,6 +394,50 @@ async def _open_agent( ) return await client.agents.create(options), True + async def _ensure_agent( + self, session: _Session, client: Any, *, resume: str | None, mcp_servers: dict[str, Any], + ) -> Literal["open", "resumed", "created"]: + """Give *session* an open agent on *client*: the one it already has, + else agent *resume* resumed, else a new one. Returns which.""" + if session.agent is not None and session.client is client: + return "open" + # None yet, or it lived on a bridge that has since been relaunched. + session.agent, is_new = await self._open_agent( + client, resume=resume, mcp_servers=mcp_servers, + ) + session.client = client + return "created" if is_new else "resumed" + + async def _send( + self, + agent: Any, + message: Any, + *, + mcp_servers: dict[str, Any], + force: bool, + max_turns: int, + on_tool: ProgressCallback | None, + ) -> TurnResult: + from cursor_sdk import LocalSendOptions, SendOptions + + usage_tracker = RunUsageTracker() + run = await agent.send( + message, + SendOptions( + # Passed on every send: a resumed agent doesn't remember its + # model, and /model takes effect here. + model=await self._model_selection(), + mcp_servers=mcp_servers, + on_delta=usage_tracker.on_delta, + local=LocalSendOptions(force=True) if force else None, + ), + ) + result = await collect_run_result( + run, max_turns=max_turns, usage_tracker=usage_tracker, on_tool=on_tool, + ) + result.session_id = agent.agent_id + return result + async def reset_session(self, session_key: str) -> None: session = self._sessions.get(session_key) if session is None: @@ -400,14 +447,6 @@ async def reset_session(self, session_key: str) -> None: # -- Model selection --------------------------------------------------- - @property - def model(self) -> str: - return self._model - - @property - def effort(self) -> str | None: - return self._effort - async def list_models(self) -> list[ModelInfo]: client = await self._get_client() models = await client.models.list(api_key=self._api_key) @@ -426,17 +465,6 @@ async def list_models(self) -> list[ModelInfo]: self._models, self._effort_params = infos, params return infos - async def _model_info(self, model_id: str) -> ModelInfo | None: - """*model_id*'s entry in the model list, or None if it isn't listed - or the list can't be fetched.""" - if not self._models: - try: - await self.list_models() - except Exception as exc: - logger.warning("Could not fetch the Cursor model list: {exc}", exc=exc) - return None - return next((m for m in self._models if m.id == model_id), None) - async def _model_selection(self) -> Any: """The model (and effort parameter, if set) to run with.""" from cursor_sdk import ModelParameterValue, ModelSelection @@ -453,24 +481,19 @@ async def _model_selection(self) -> Any: params=[ModelParameterValue(id=param, value=self._effort)], ) - async def effort_levels(self) -> list[str] | None: - info = await self._model_info(self._model) - return list(info.effort_levels) if info else None - async def set_model(self, model_id: str) -> str: - model_id = model_id.strip() - info = await self._model_info(model_id) - if info is None and self._models: - raise ValueError(f"Unknown model '{model_id}'.") + """Switch from the next turn on. The model is passed with every send, + so open agents keep their conversation and need no reconnect.""" + info = await self._known_model(model_id.strip()) effort = self._effort note = "" - if info is not None and self._effort and self._effort not in info.effort_levels: + if self._effort and self._effort not in info.effort_levels: effort = None note = ( f"{info.display_name} doesn't support '{self._effort}' effort, " "so its default is used." ) - self._apply_model(model_id, effort) + self._apply_model(info.id, effort) return note async def set_effort(self, level: str) -> None: @@ -488,24 +511,6 @@ async def set_effort(self, level: str) -> None: ) self._apply_model(self._model, match) - def _apply_model(self, model: str, effort: str | None) -> None: - """Use *model* / *effort* from the next turn on and save them to - ~/.memclaw/.env. The model is passed with every send, so open agents - keep their conversation and need no reconnect.""" - from ..setup import update_env_file # local import — setup imports backends - - logger.info( - "Cursor model: {m} effort: {e} -> {m2} effort: {e2}", - m=self._model, e=self._effort, m2=model, e2=effort, - ) - self._model, self._effort = model, effort - self.config.cursor_model = model - self.config.cursor_effort = effort or "" - if effort: - update_env_file({"CURSOR_MODEL": model, "CURSOR_EFFORT": effort}) - else: - update_env_file({"CURSOR_MODEL": model}, drop_keys=["CURSOR_EFFORT"]) - # -- Runtime ----------------------------------------------------------- async def run_one_shot(self, *, system_prompt: str, user_message: str) -> str: @@ -552,7 +557,7 @@ async def run_turn( max_turns: int = 10, on_tool: ProgressCallback | None = None, ) -> TurnResult: - from cursor_sdk import CursorAgentError, LocalSendOptions, SendOptions + from cursor_sdk import CursorAgentError if not self._api_key: raise RuntimeError("CURSOR_API_KEY is not configured") @@ -566,44 +571,26 @@ async def run_turn( session = self._sessions.setdefault(session_key, _Session()) async with session.lock: client = await self._get_client() - if session.agent is not None and session.client is not client: - session.agent = None # its bridge was relaunched - opened = is_new = False try: - if session.agent is None: - session.agent, is_new = await self._open_agent( - client, resume=resume_session_id, mcp_servers=mcp_servers, - ) - session.client = client - opened = True + agent = await self._ensure_agent( + session, client, resume=resume_session_id, mcp_servers=mcp_servers, + ) message = _build_user_message( - system_prompt=system_prompt if is_new else "", + # A Cursor agent has no system-prompt field; a new one + # gets it at the top of its first message. + system_prompt=system_prompt if agent == "created" else "", context=context, user_message=user_message, image_b64=image_b64, image_media_type=image_media_type, ) - usage_tracker = RunUsageTracker() - run = await session.agent.send( - message, - SendOptions( - # Passed on every send: a resumed agent doesn't - # remember its model, and /model takes effect here. - model=await self._model_selection(), - mcp_servers=mcp_servers, - on_delta=usage_tracker.on_delta, - # A resumed agent may still show a run that died with - # the previous process; take it over. - local=LocalSendOptions(force=True) if opened and not is_new else None, - ), - ) - result = await collect_run_result( - run, - max_turns=max_turns, - usage_tracker=usage_tracker, - on_tool=on_tool, + result = await self._send( + session.agent, message, mcp_servers=mcp_servers, + # A resumed agent may still show a run that died with + # the previous process; take it over. + force=agent == "resumed", + max_turns=max_turns, on_tool=on_tool, ) - result.session_id = session.agent.agent_id except BaseException as exc: # The agent may be mid-run or gone; reopen (resuming) next time. await self._close_agent(session) diff --git a/memclaw/backends/model_selection.py b/memclaw/backends/model_selection.py new file mode 100644 index 0000000..aef1bf9 --- /dev/null +++ b/memclaw/backends/model_selection.py @@ -0,0 +1,84 @@ +"""/model and /effort state shared by the backends that list their models.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar + +from loguru import logger + +if TYPE_CHECKING: + from ..config import MemclawConfig + from .claude_models import ModelInfo + + +class ModelSelectionMixin: + """The current model and effort, lookups in the backend's model list, + and saving a switch to ~/.memclaw/.env. + + A backend sets `_label` (for logs) and `_env_prefix` (CLAUDE → the + CLAUDE_MODEL / CLAUDE_EFFORT keys and the claude_model / claude_effort + config fields), keeps `_model`, `_effort` and `_models`, and implements + `list_models()`, `set_model()` and `set_effort()`. + """ + + _label: ClassVar[str] + _env_prefix: ClassVar[str] + + config: "MemclawConfig" + _model: str + _effort: str | None + _models: list["ModelInfo"] + + async def list_models(self) -> list["ModelInfo"]: + raise NotImplementedError + + @property + def model(self) -> str: + return self._model + + @property + def effort(self) -> str | None: + return self._effort + + async def _model_info(self, model_id: str) -> "ModelInfo | None": + """*model_id*'s entry in the model list, or None if it isn't listed + or the list can't be fetched.""" + if not self._models: + try: + await self.list_models() + except Exception as exc: + logger.warning("Could not fetch the {b} model list: {exc}", b=self._label, exc=exc) + return None + return next((m for m in self._models if m.id == model_id), None) + + async def _known_model(self, model_id: str) -> "ModelInfo": + """*model_id*'s entry, or ValueError. An id that can't be checked + (the list didn't load) is refused too: a typo would break every turn + until ~/.memclaw/.env is fixed by hand.""" + info = await self._model_info(model_id) + if info is not None: + return info + if not self._models: + raise ValueError("Couldn't load the model list; try again later.") + raise ValueError(f"Unknown model '{model_id}'.") + + async def effort_levels(self) -> list[str] | None: + info = await self._model_info(self._model) + return list(info.effort_levels) if info else None + + def _apply_model(self, model: str, effort: str | None) -> None: + """Use *model* / *effort* from the next turn on and save them.""" + from ..setup import update_env_file # local import — setup imports backends + + logger.info( + "{b} model: {m} effort: {e} -> {m2} effort: {e2}", + b=self._label, m=self._model, e=self._effort, m2=model, e2=effort, + ) + self._model, self._effort = model, effort + prefix = self._env_prefix + setattr(self.config, f"{prefix.lower()}_model", model) + setattr(self.config, f"{prefix.lower()}_effort", effort or "") + if effort: + update_env_file({f"{prefix}_MODEL": model, f"{prefix}_EFFORT": effort}) + else: + update_env_file({f"{prefix}_MODEL": model}, drop_keys=[f"{prefix}_EFFORT"]) diff --git a/tests/backends/test_cursor.py b/tests/backends/test_cursor.py index b738d6e..dd70d8a 100644 --- a/tests/backends/test_cursor.py +++ b/tests/backends/test_cursor.py @@ -782,6 +782,17 @@ async def test_set_unknown_model_raises(self, tmp_path): with _patched(backend, client), pytest.raises(ValueError, match="Unknown model"): await backend.set_model("nope") + @pytest.mark.asyncio + async def test_switch_is_refused_when_the_list_is_unavailable(self, tmp_path): + backend = _backend(tmp_path) + client = _make_client() + client.models.list.side_effect = RuntimeError("offline") + + with _patched(backend, client), pytest.raises(ValueError, match="Couldn't load"): + await backend.set_model("composr-2.5") + assert backend.model == "composer-2.5" + assert "CURSOR_MODEL" not in self._env() + @pytest.mark.asyncio async def test_effort_is_sent_as_model_parameter(self, tmp_path): backend = _backend(tmp_path) From deefab862f12decf5ced053b0bf51c32a894e326 Mon Sep 17 00:00:00 2001 From: Emanuilo Jovanovic Date: Sun, 18 Oct 2026 19:11:26 +0200 Subject: [PATCH 09/10] build: pin cursor-sdk to 0.1.x The Cursor backend is written and checked against the 0.1.x SDK; CI installed 1.0.38 (>=0.1.0 allowed it), which no longer exports UnknownAgentError. Pin to <0.2 so a major SDK upgrade is a deliberate change. --- pyproject.toml | 4 ++-- uv.lock | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index fbeecc7..e3dfa05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,13 +51,13 @@ Issues = "https://github.com/emanuilo/memclaw/issues" [project.optional-dependencies] cursor = [ - "cursor-sdk>=0.1.0", + "cursor-sdk>=0.1.4,<0.2", ] dev = [ "pytest>=9.0.0", "pytest-asyncio>=1.0.0", "ruff>=0.9.0", - "cursor-sdk>=0.1.0", + "cursor-sdk>=0.1.4,<0.2", ] [tool.hatch.build.targets.wheel] diff --git a/uv.lock b/uv.lock index c999ef0..cbf9c90 100644 --- a/uv.lock +++ b/uv.lock @@ -1007,8 +1007,8 @@ requires-dist = [ { name = "claude-agent-sdk", specifier = ">=0.2.163,<0.3" }, { name = "click", specifier = ">=8.0.0" }, { name = "curl-cffi", specifier = ">=0.7.0" }, - { name = "cursor-sdk", marker = "extra == 'cursor'", specifier = ">=0.1.0" }, - { name = "cursor-sdk", marker = "extra == 'dev'", specifier = ">=0.1.0" }, + { name = "cursor-sdk", marker = "extra == 'cursor'", specifier = ">=0.1.4,<0.2" }, + { name = "cursor-sdk", marker = "extra == 'dev'", specifier = ">=0.1.4,<0.2" }, { name = "httpx", specifier = ">=0.27.0" }, { name = "loguru", specifier = ">=0.7.0" }, { name = "mcp", specifier = ">=1.23.0,<2" }, From 2f18a303f6464e92d5beb0c790105fa5831ccc09 Mon Sep 17 00:00:00 2001 From: Emanuilo Jovanovic Date: Sun, 18 Oct 2026 19:14:14 +0200 Subject: [PATCH 10/10] test: cursor-sdk's not-found error under both names 0.1.4 calls it UnknownAgentError, 0.1.9 AgentNotFoundError (both CursorAgentError subclasses, which is what the backend catches). The suite passes on both versions. --- tests/backends/test_cursor.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/backends/test_cursor.py b/tests/backends/test_cursor.py index dd70d8a..5604b9b 100644 --- a/tests/backends/test_cursor.py +++ b/tests/backends/test_cursor.py @@ -505,12 +505,15 @@ async def test_resume_passes_model_and_forces_first_send(self, tmp_path): @pytest.mark.asyncio async def test_resume_failure_starts_fresh_agent(self, tmp_path): - from cursor_sdk import UnknownAgentError + import cursor_sdk + # The SDK's "no such agent" error: UnknownAgentError up to 0.1.4, + # AgentNotFoundError later. Both are CursorAgentErrors. + not_found = getattr(cursor_sdk, "AgentNotFoundError", None) or cursor_sdk.UnknownAgentError backend = _backend(tmp_path) fresh = _make_agent("agent-new") client = _make_client( - resume=[UnknownAgentError("Agent agent-gone not found")], create=[fresh], + resume=[not_found("Agent agent-gone not found")], create=[fresh], ) with _patched(backend, client):