diff --git a/AGENTS.md b/AGENTS.md index 4953bd66..f1bc7da7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -98,7 +98,7 @@ Place state according to its lifetime: Application restart. Use narrow, namespaced keys in `app_context.metadata`, following existing patterns such as -`tool_contexts` and `channel_sink`. The ApplicationContext is shared, so account for +`tool_contexts`. The ApplicationContext is shared, so account for concurrent access when values are mutable. New Step code must not fall back to `self.kwargs` or another Step field to emulate shared state when `app_context` is absent; tests of shared state should construct an `ApplicationContext`. If shared state grows into a stable diff --git a/docs/en/framework.md b/docs/en/framework.md index b4bc11a7..ec3836c1 100644 --- a/docs/en/framework.md +++ b/docs/en/framework.md @@ -139,10 +139,14 @@ Configuration parsing supports: `BaseService.run_app()` executes in this order: +Set the optional `service.jobs` list to restrict HTTP or MCP exposure to those job names. If omitted, all jobs with +`enable_serve: true` remain eligible; an empty list exposes none. The whitelist does not override `enable_serve: false`. +When the list is configured, a missing, disabled, unsupported, or invalid selected job fails service startup. + ```mermaid flowchart LR A["Service.build_service(app)"] --> B["read app.context.jobs"] - B --> C{"job.enable_serve == true?"} + B --> C{"enabled and selected by service.jobs?"} C -->|yes| D["Service.add_job(job)"] C -->|no| E["skip registration"] D --> F["Service.start_service(app)"] @@ -167,6 +171,9 @@ MCP service behavior: | `StreamJob` | Currently skipped and not registered. | | `BackgroundJob` | Forces `enable_serve=False` at construction and is never exposed. | +MCP services can inject server-owned arguments with `injected_job_kwargs`; callers cannot override those arguments. +Set `tool_error_on_failure: true` to expose an unsuccessful ReMe `Response` as an MCP tool error. + ## 4. Registry and Dependency Injection ### 4.1 Global Registry R diff --git a/docs/zh/framework.md b/docs/zh/framework.md index 6d4af663..b67cbffd 100644 --- a/docs/zh/framework.md +++ b/docs/zh/framework.md @@ -134,10 +134,14 @@ reme search query="memory" backend=mcp `BaseService.run_app()` 的顺序： +可通过可选的 `service.jobs` 列表将 HTTP 或 MCP 仅暴露给指定 Job。未配置时，所有 `enable_serve: true` 的 Job +仍可被暴露；配置为空列表时不暴露任何 Job。该白名单不会覆盖 `enable_serve: false`。 +配置该列表后，缺失、禁用、不受支持或无效的已选 Job 会导致服务启动失败。 + ```mermaid flowchart LR A["Service.build_service(app)"] --> B["读取 app.context.jobs"] - B --> C{"job.enable_serve == true?"} + B --> C{"已启用且被 service.jobs 选中？"} C -->|是| D["Service.add_job(job)"] C -->|否| E["跳过注册"] D --> F["Service.start_service(app)"] @@ -162,6 +166,9 @@ MCP service 行为： | `StreamJob` | 当前跳过，不注册 | | `BackgroundJob` | 构造时强制 `enable_serve=False`，不会暴露 | +MCP 服务可通过 `injected_job_kwargs` 注入由服务端管理的参数，调用方不能覆盖这些参数。设置 +`tool_error_on_failure: true` 后，不成功的 ReMe `Response` 会作为 MCP tool error 返回。 + ## 4. Registry 与依赖注入 ### 4.1 全局注册表 R diff --git a/pyproject.toml b/pyproject.toml index 4ccdb2c6..573b098b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,14 +42,14 @@ dependencies = [ [project.optional-dependencies] core = [ - "agentscope==2.0.4", - "claude-agent-sdk>=0.2.91", + "agentscope==2.0.4.post1", + "claude-agent-sdk>=0.2.123", "faiss-cpu>=1.13.2", "jieba>=0.42.1", "rjieba>=0.2.1", "neo4j>=6.2.0", "networkx>=3.4.2", - "openai-codex>=0.1.0b3", + "openai-codex>=0.144.4", ] dev = [ "pre-commit", diff --git a/reme/__init__.py b/reme/__init__.py index 167c3a74..d0a37a80 100644 --- a/reme/__init__.py +++ b/reme/__init__.py @@ -1,6 +1,6 @@ """ReMe CLI package.""" -__version__ = "0.4.1.3" +__version__ = "0.4.1.4" from . import config from . import constants diff --git a/reme/components/agent_wrapper/__init__.py b/reme/components/agent_wrapper/__init__.py index f5c52115..7a0ce24c 100644 --- a/reme/components/agent_wrapper/__init__.py +++ b/reme/components/agent_wrapper/__init__.py @@ -2,7 +2,8 @@ from .base_agent_wrapper import BaseAgentWrapper from .as_agent_wrapper import AsAgentWrapper -from .cc_agent_wrapper import CcAgentWrapper, CcFileSessionStore +from .cc_agent_wrapper import CcAgentWrapper +from .cc_session_store import CcFileSessionStore from .codex_agent_wrapper import CodexAgentWrapper __all__ = ["BaseAgentWrapper", "AsAgentWrapper", "CcAgentWrapper", "CcFileSessionStore", "CodexAgentWrapper"] diff --git a/reme/components/agent_wrapper/as_agent_wrapper.py b/reme/components/agent_wrapper/as_agent_wrapper.py index 658fcfd0..4638987e 100644 --- a/reme/components/agent_wrapper/as_agent_wrapper.py +++ b/reme/components/agent_wrapper/as_agent_wrapper.py @@ -57,7 +57,6 @@ from ..component_registry import R from ...enumeration import ChunkEnum from ...schema import StreamChunk from ...utils import AsStateHandler -from ...utils.env_utils import load_env if TYPE_CHECKING: from ..job.base_job import BaseJob @@ -110,6 +109,8 @@ class BypassAnalysisBash(Bash): class AsAgentWrapper(BaseAgentWrapper): """Agent wrapper backed by AgentScope framework.""" + SDK_PACKAGE = "agentscope" + def __init__(self, as_llm: str = "default", session_retention_days: int = 10, **kwargs): super().__init__(**kwargs) self.as_llm = self.bind(as_llm, BaseAsLLM, optional=False) @@ -256,11 +257,6 @@ class AsAgentWrapper(BaseAgentWrapper): skills = [skills] return [str(self.project_skills_root / skill) for skill in skills] - def _load_tool_env(self) -> dict[str, str]: - """Load project environment variables for tools spawned by AgentScope.""" - project_env = self.project_path / ".env" - return load_env(project_env) if project_env.exists() else load_env() - async def _build_agent(self, inputs: Any, **kwargs) -> tuple[Agent, Any]: """Build an Agent instance from kwargs. Returns (agent, processed_inputs).""" model = self.as_llm.model if self.as_llm else None @@ -268,7 +264,6 @@ class AsAgentWrapper(BaseAgentWrapper): raise ValueError("AsAgentWrapper requires a bound as_llm component with a valid model.") self._cleanup_expired_sessions() - self._load_tool_env() system_prompt = kwargs.get("system_prompt", "You are a helpful assistant.") job_tools: list[str] = kwargs.get("job_tools", []) diff --git a/reme/components/agent_wrapper/base_agent_wrapper.py b/reme/components/agent_wrapper/base_agent_wrapper.py index 9267f05e..4287eda6 100644 --- a/reme/components/agent_wrapper/base_agent_wrapper.py +++ b/reme/components/agent_wrapper/base_agent_wrapper.py @@ -2,8 +2,9 @@ from abc import abstractmethod from collections.abc import AsyncGenerator +from importlib import metadata from pathlib import Path -from typing import Any, TYPE_CHECKING +from typing import Any, ClassVar, TYPE_CHECKING from pydantic import BaseModel @@ -19,10 +20,17 @@ class BaseAgentWrapper(BaseComponent): """Abstract base for agent wrapper components with swappable backends.""" component_type = ComponentEnum.AGENT_WRAPPER + SDK_PACKAGE: ClassVar[str | None] = None def __init__(self, cwd: str | Path | None = None, **kwargs) -> None: super().__init__(**kwargs) self._cwd = cwd + if self.SDK_PACKAGE: + try: + sdk_version = metadata.version(self.SDK_PACKAGE) + except metadata.PackageNotFoundError: + sdk_version = "unknown" + self.logger.info(f"Agent SDK package={self.SDK_PACKAGE} version={sdk_version}") @property def cwd(self) -> Path: @@ -62,6 +70,13 @@ class BaseAgentWrapper(BaseComponent): """Project-level skills directory shared by agent backends.""" return self.project_path / "skills" + @property + def subprocess_environment(self) -> dict[str, str]: + """Configured environment variables for child agent processes.""" + if self.app_context is None: + return {} + return self.app_context.app_config.environment + def set_output_schema(self, schema: dict | type[BaseModel]) -> "BaseAgentWrapper": """Set a JSON schema for structured output. Accepts dict or BaseModel class. Returns self for chaining.""" self.kwargs["output_schema"] = self._normalize_output_schema(schema) diff --git a/reme/components/agent_wrapper/cc_agent_wrapper.py b/reme/components/agent_wrapper/cc_agent_wrapper.py index 7402c06d..66903272 100644 --- a/reme/components/agent_wrapper/cc_agent_wrapper.py +++ b/reme/components/agent_wrapper/cc_agent_wrapper.py @@ -1,9 +1,9 @@ """Claude Code SDK backend for the unified agent wrapper.""" import json -import os from collections.abc import AsyncGenerator -from dataclasses import asdict +from contextlib import aclosing +from dataclasses import asdict, dataclass, fields from pathlib import Path from typing import Any, TYPE_CHECKING @@ -11,200 +11,29 @@ from .base_agent_wrapper import BaseAgentWrapper from ..component_registry import R from ...enumeration import ChunkEnum from ...schema import StreamChunk -from ...utils.env_utils import load_env if TYPE_CHECKING: + from claude_agent_sdk import AssistantMessage, ResultMessage, UserMessage + from ..job.base_job import BaseJob - from claude_agent_sdk.types import SessionKey, SessionStoreEntry, SessionStoreListEntry -class CcFileSessionStore: - """File-backed Claude Code SessionStore rooted under a given directory. +@dataclass(frozen=True) +class _BlockState: + """Metadata needed to correlate one streamed content block.""" - Rooted under the ReMe workspace for the inner agent's own sessions, but the - same reader is reused (with a different root) to load an *outer* Claude Code - session's transcript by id — see ``AutoMemoryCCStep``. - """ - - def __init__(self, root: Path) -> None: - self.root = root - - @staticmethod - def _safe_parts(value: str) -> list[str]: - parts = [part for part in value.split("/") if part] - if not parts or any(part in {".", ".."} for part in parts): - raise ValueError(f"Invalid session store path component: {value!r}") - return parts - - def _path_for_key(self, key: "SessionKey") -> Path: - session_id = key["session_id"] - subpath = key.get("subpath") - - path = self.root.joinpath(*self._safe_parts(session_id)) - if subpath: - path = path.joinpath(*self._safe_parts(subpath)) - else: - path = path.with_suffix(".jsonl") - if subpath: - path = path.with_suffix(".jsonl") - - resolved_root = self.root.resolve() - resolved_path = path.resolve() - if resolved_root != resolved_path and resolved_root not in resolved_path.parents: - raise ValueError(f"Session store path escapes root: {resolved_path}") - return path - - @staticmethod - def _read_entries(path: Path) -> list["SessionStoreEntry"]: - """Read JSONL session-store entries from disk.""" - if not path.exists(): - return [] - entries = [] - for line in path.read_text(encoding="utf-8").splitlines(): - if line.strip(): - entries.append(json.loads(line)) - return entries - - async def append(self, key: "SessionKey", entries: list["SessionStoreEntry"]) -> None: - """Append new session-store entries, deduplicating by UUID.""" - path = self._path_for_key(key) - path.parent.mkdir(parents=True, exist_ok=True) - - existing_uuids = { - entry.get("uuid") for entry in self._read_entries(path) if isinstance(entry, dict) and entry.get("uuid") - } - new_entries = [ - entry for entry in entries if not (isinstance(entry, dict) and entry.get("uuid") in existing_uuids) - ] - if not new_entries: - return - - with path.open("a", encoding="utf-8") as f: - for entry in new_entries: - f.write(json.dumps(entry, ensure_ascii=False, separators=(",", ":")) + "\n") - - async def load(self, key: "SessionKey") -> list["SessionStoreEntry"] | None: - """Load session-store entries for a key.""" - path = self._path_for_key(key) - if not path.exists(): - return None - return self._read_entries(path) - - async def list_sessions(self, _project_key: str) -> list["SessionStoreListEntry"]: - """List root-level Claude Code sessions.""" - if not self.root.exists(): - return [] - return [ - {"session_id": path.stem, "mtime": int(path.stat().st_mtime * 1000)} - for path in self.root.glob("*.jsonl") - if path.is_file() - ] - - async def delete(self, key: "SessionKey") -> None: - """Delete a session-store entry and any subkey directory.""" - path = self._path_for_key(key) - if path.exists(): - path.unlink() - - if not key.get("subpath"): - session_dir = self.root.joinpath(*self._safe_parts(key["session_id"])) - if session_dir.exists(): - for child in sorted(session_dir.rglob("*"), reverse=True): - if child.is_file(): - child.unlink() - elif child.is_dir(): - child.rmdir() - session_dir.rmdir() - - async def list_subkeys(self, key: dict[str, str]) -> list[str]: - """List subkeys below a root session key.""" - session_dir = self.root.joinpath(*self._safe_parts(key["session_id"])) - if not session_dir.exists(): - return [] - subkeys = [] - for path in session_dir.rglob("*.jsonl"): - if path.is_file(): - subkeys.append(str(path.relative_to(session_dir).with_suffix(""))) - return subkeys + block_id: str | None + block_type: str + tool_name: str | None @R.register("claude_code") class CcAgentWrapper(BaseAgentWrapper): """Agent wrapper backed by Claude Code SDK.""" + SDK_PACKAGE = "claude-agent-sdk" DEFAULT_DISALLOWED_TOOLS = ["WebSearch"] - SYSTEM_PROMPT_MODES = {"append", "replace"} - - @staticmethod - def _first_non_empty(*values: Any) -> str: - for value in values: - if isinstance(value, str) and value: - return value - return "" - - def _default_llm_credential(self) -> dict[str, Any]: - """Return the default as_llm credential config, if available.""" - if self.app_context is None: - return {} - components = self.app_context.app_config.components - llm_configs = components.get("as_llm") or components.get("AS_LLM") or components.get("as_llm".upper()) - if llm_configs is None: - from ...enumeration import ComponentEnum - - llm_configs = components.get(ComponentEnum.AS_LLM) - if not isinstance(llm_configs, dict): - return {} - - default_llm = llm_configs.get("default") - credential = getattr(default_llm, "credential", None) - return credential if isinstance(credential, dict) else {} - - def _claude_code_api_env(self, kwargs: dict[str, Any]) -> dict[str, str]: - """Resolve Anthropic-compatible API environment for Claude Code.""" - credential = kwargs.get("credential") if isinstance(kwargs.get("credential"), dict) else {} - default_credential = self._default_llm_credential() - - base_url = self._first_non_empty( - kwargs.get("base_url"), - credential.get("base_url"), - os.getenv("ANTHROPIC_BASE_URL"), - os.getenv("CLAUDE_CODE_BASE_URL"), - os.getenv("LLM_BASE_URL"), - default_credential.get("base_url"), - ) - api_key = self._first_non_empty( - kwargs.get("api_key"), - credential.get("api_key"), - os.getenv("ANTHROPIC_AUTH_TOKEN"), - os.getenv("CLAUDE_CODE_API_KEY"), - os.getenv("LLM_API_KEY"), - default_credential.get("api_key"), - ) - - env: dict[str, str] = {} - if base_url: - env["ANTHROPIC_BASE_URL"] = base_url - if api_key: - env["ANTHROPIC_AUTH_TOKEN"] = api_key - return env - - @classmethod - def _apply_system_prompt_mode(cls, kwargs: dict[str, Any]) -> None: - """Translate the configured prompt mode into Claude SDK semantics.""" - mode = kwargs.pop("system_prompt_mode", "replace") - if mode not in cls.SYSTEM_PROMPT_MODES: - allowed = ", ".join(sorted(cls.SYSTEM_PROMPT_MODES)) - raise ValueError(f"Unknown system_prompt_mode {mode!r}; expected one of: {allowed}") - - if mode == "append" and "system_prompt" in kwargs: - prompt = kwargs["system_prompt"] - if not isinstance(prompt, str): - raise TypeError("system_prompt must be a string when system_prompt_mode='append'") - kwargs["system_prompt"] = { - "type": "preset", - "preset": "claude_code", - "append": prompt, - } + MCP_SERVER_NAME = "mcp_server" @property def session_path(self) -> Path: @@ -238,14 +67,10 @@ class CcAgentWrapper(BaseAgentWrapper): for target in (self.project_path / ".claude" / "skills", config_dir / "skills"): try: - # Migrate directory-level links created by older ReMe versions. if target.is_symlink(): - if target.resolve() == project_skills.resolve(): - target.unlink() - else: - self.logger.warning(f"Preserving existing Claude Code skills link: {target}") - continue - elif target.exists() and not target.is_dir(): + self.logger.warning(f"Preserving existing Claude Code skills link: {target}") + continue + if target.exists() and not target.is_dir(): self.logger.warning(f"Preserving existing Claude Code skills path: {target}") continue @@ -268,13 +93,23 @@ class CcAgentWrapper(BaseAgentWrapper): from claude_agent_sdk import SdkMcpTool async def run_job(args): + call_args = dict(args) if tool_context_id: - assert "tool_context_id" not in args, "tool_context_id is injected by agent_wrapper" - args["tool_context_id"] = tool_context_id - response = await job(**args) - return {"content": [{"type": "text", "text": str(response.answer)}], "is_error": not response.success} + if "tool_context_id" in call_args: + raise ValueError("tool_context_id is injected by agent_wrapper") + call_args["tool_context_id"] = tool_context_id + response = await job(**call_args) + return { + "content": [{"type": "text", "text": str(response.answer)}], + "is_error": not response.success, + } - return SdkMcpTool(name=job.name, description=job.description, input_schema=job.parameters, handler=run_job) + return SdkMcpTool( + name=job.name, + description=job.description, + input_schema=job.parameters, + handler=run_job, + ) def _build_options(self, inputs: Any, stream: bool = False, **kwargs) -> Any: """Build ClaudeAgentOptions from kwargs. @@ -283,43 +118,49 @@ class CcAgentWrapper(BaseAgentWrapper): ``StreamEvent`` messages are emitted alongside the final ``ResultMessage``. """ - from claude_agent_sdk import create_sdk_mcp_server - from claude_agent_sdk.types import ClaudeAgentOptions + from claude_agent_sdk import ClaudeAgentOptions, create_sdk_mcp_server - self._apply_system_prompt_mode(kwargs) + if not isinstance(inputs, str): + raise NotImplementedError("Only string input is supported for Claude Code.") - skills = kwargs.get("skills") - if isinstance(skills, str) and skills != "all": - kwargs["skills"] = [skills] + selected_skills = kwargs.get("skills") + if isinstance(selected_skills, str) and selected_skills != "all": + selected_skills = [selected_skills] if "setting_sources" not in kwargs and kwargs.get("skills") is None: kwargs["setting_sources"] = [] - disallowed_tools = list(kwargs.get("disallowed_tools") or []) - for tool_name in self.DEFAULT_DISALLOWED_TOOLS: - if tool_name not in disallowed_tools: - disallowed_tools.append(tool_name) - kwargs["disallowed_tools"] = disallowed_tools - - opts = ClaudeAgentOptions() - if stream: - opts.include_partial_messages = True - skip_keys = {"job_tools", "output_schema", "api_key", "base_url", "credential"} - for k, v in kwargs.items(): - if k not in skip_keys and hasattr(opts, k): - setattr(opts, k, v) + option_fields = {field.name for field in fields(ClaudeAgentOptions)} + option_kwargs = {key: value for key, value in kwargs.items() if key not in skip_keys and key in option_fields} + option_kwargs["disallowed_tools"] = list( + dict.fromkeys( + [ + *(kwargs.get("disallowed_tools") or []), + *self.DEFAULT_DISALLOWED_TOOLS, + ], + ), + ) + if selected_skills is not None: + option_kwargs["skills"] = selected_skills + if stream: + option_kwargs["include_partial_messages"] = True + opts = ClaudeAgentOptions(**option_kwargs) - model = getattr(opts, "model", None) or kwargs.get("model") - project_env = self.project_path / ".env" - opts.env.update(load_env(project_env) if project_env.exists() else load_env()) - extra_env_dict: dict = self._claude_code_api_env(kwargs) - if model: + opts.env = dict(opts.env) + opts.env.update(self.subprocess_environment) + api_key = kwargs.get("api_key") + base_url = kwargs.get("base_url") + extra_env_dict = { + "ANTHROPIC_AUTH_TOKEN": api_key if isinstance(api_key, str) else "", + "ANTHROPIC_BASE_URL": base_url if isinstance(base_url, str) else "", + } + if opts.model: extra_env_dict.update( { - "ANTHROPIC_MODEL": model, - "ANTHROPIC_DEFAULT_HAIKU_MODEL": model, - "ANTHROPIC_DEFAULT_SONNET_MODEL": model, - "ANTHROPIC_DEFAULT_OPUS_MODEL": model, + "ANTHROPIC_MODEL": opts.model, + "ANTHROPIC_DEFAULT_HAIKU_MODEL": opts.model, + "ANTHROPIC_DEFAULT_SONNET_MODEL": opts.model, + "ANTHROPIC_DEFAULT_OPUS_MODEL": opts.model, }, ) opts.env.update(extra_env_dict) @@ -327,25 +168,28 @@ class CcAgentWrapper(BaseAgentWrapper): opts.cwd = opts.cwd or self.cwd claude_config_dir = self.session_path / "claude_config" opts.env.setdefault("CLAUDE_CONFIG_DIR", str(claude_config_dir)) - if opts.skills is not None: - self._ensure_claude_skill_dir(claude_config_dir, opts.skills) - opts.session_store = opts.session_store or CcFileSessionStore(self.session_path / "claude_code") + if selected_skills is not None: + self._ensure_claude_skill_dir(claude_config_dir, selected_skills) job_tools: list[str] = kwargs.get("job_tools", []) resolved_jobs = self._resolve_job_tools(job_tools) if resolved_jobs: + if not isinstance(opts.mcp_servers, dict): + raise ValueError("job_tools require mcp_servers to be a mapping so the ReMe SDK server can be merged") + opts.mcp_servers = dict(opts.mcp_servers) + if self.MCP_SERVER_NAME in opts.mcp_servers: + raise ValueError(f"mcp_servers already contains reserved server name {self.MCP_SERVER_NAME!r}") sdk_tools = [self._make_tool(job, kwargs.get("tool_context_id")) for job in resolved_jobs] - server = create_sdk_mcp_server(name="mcp_server", tools=sdk_tools) - opts.mcp_servers = opts.mcp_servers if isinstance(opts.mcp_servers, dict) else {} - opts.mcp_servers["mcp_server"] = server + opts.mcp_servers[self.MCP_SERVER_NAME] = create_sdk_mcp_server( + name=self.MCP_SERVER_NAME, + tools=sdk_tools, + ) + opts.allowed_tools = list(opts.allowed_tools) opts.allowed_tools.extend(job.name for job in resolved_jobs) if (output_schema := kwargs.get("output_schema")) is not None: opts.output_format = {"type": "json_schema", "schema": output_schema} - if not isinstance(inputs, str): - raise NotImplementedError("Only string input is supported for Claude Code.") - return opts # ----- StreamChunk conversion ------------------------------------------- @@ -356,17 +200,12 @@ class CcAgentWrapper(BaseAgentWrapper): cls, raw: dict, session_id: str | None = None, - block_ids: dict[int, str] | None = None, - block_types: dict[int, str] | None = None, - tool_call_names: dict[int, str] | None = None, + block_states: dict[int, _BlockState] | None = None, ) -> StreamChunk | None: """Convert a raw Anthropic streaming event dict to a StreamChunk. - ``block_ids`` / ``block_types`` / ``tool_call_names`` map - content-block ``index`` to metadata tracked from the - ``content_block_start`` event, so that later delta / stop - events can reference the correct ``block_id`` and - ``chunk_type``. + ``block_states`` maps each content-block index to metadata captured at + ``content_block_start`` for use by later delta and stop events. Returns ``None`` for events that should be silently skipped. """ @@ -375,17 +214,23 @@ class CcAgentWrapper(BaseAgentWrapper): # --- Message-level lifecycle ---------------------------------------- if event_type == "message_start": + if block_states is not None: + block_states.clear() message = raw.get("message", {}) - meta = {"message_id": message.get("id"), "model": message.get("model"), "role": message.get("role")} + meta = { + "message_id": message.get("id"), + "model": message.get("model"), + "role": message.get("role"), + } return cls._chunk(ChunkEnum.REPLY_START, session_id=session_id, chunk="", metadata=meta) if event_type == "message_delta": delta = raw.get("delta", {}) usage = raw.get("usage", {}) return cls._chunk( - ChunkEnum.REPLY_END, + ChunkEnum.USAGE, session_id=session_id, - chunk="", + chunk=json.dumps(usage), output_tokens=usage.get("output_tokens"), metadata={"stop_reason": delta.get("stop_reason")}, ) @@ -398,21 +243,22 @@ class CcAgentWrapper(BaseAgentWrapper): if event_type == "content_block_start": idx, content_block = raw.get("index", 0), raw.get("content_block", {}) block_type, bid = content_block.get("type", ""), content_block.get("id", "") - - # Track for later delta / stop correlation - if block_ids is not None and bid: - block_ids[idx] = bid - if block_types is not None and block_type: - block_types[idx] = block_type - if tool_call_names is not None and content_block.get("name"): - tool_call_names[idx] = content_block["name"] + if block_states is not None: + block_states[idx] = _BlockState(bid or None, block_type, content_block.get("name")) if block_type == "text": return cls._chunk(ChunkEnum.CONTENT, block_id=bid, chunk=content_block.get("text", "")) if block_type == "thinking": - return cls._chunk(ChunkEnum.THINK, block_id=bid, chunk=content_block.get("thinking", "")) - if block_type == "tool_use": - payload = {"name": content_block.get("name"), "id": content_block.get("id")} + return cls._chunk( + ChunkEnum.THINK, + block_id=bid, + chunk=content_block.get("thinking", ""), + ) + if block_type in {"tool_use", "server_tool_use"}: + payload = { + "name": content_block.get("name"), + "id": content_block.get("id"), + } return cls._chunk( ChunkEnum.TOOL_CALL, block_id=bid, @@ -426,8 +272,9 @@ class CcAgentWrapper(BaseAgentWrapper): delta = raw.get("delta", {}) delta_type = delta.get("type", "") idx = raw.get("index", 0) - bid = block_ids.get(idx) if block_ids else None - tc_name = tool_call_names.get(idx) if tool_call_names else None + state = block_states.get(idx) if block_states else None + bid = state.block_id if state else None + tool_name = state.tool_name if state else None if delta_type == "text_delta": return cls._chunk(ChunkEnum.CONTENT, block_id=bid, chunk=delta.get("text", "")) @@ -438,20 +285,27 @@ class CcAgentWrapper(BaseAgentWrapper): ChunkEnum.TOOL_CALL, block_id=bid, tool_call_id=bid, - tool_call_name=tc_name, + tool_call_name=tool_name, chunk=delta.get("partial_json", ""), ) return None if event_type == "content_block_stop": idx = raw.get("index", 0) - bid = block_ids.get(idx) if block_ids else None - btype = block_types.get(idx) if block_types else None - tc_name = tool_call_names.get(idx) if tool_call_names else None + state = block_states.pop(idx, None) if block_states else None + bid = state.block_id if state else None + block_type = state.block_type if state else None + tool_name = state.tool_name if state else None - if btype == "tool_use": - return cls._chunk(ChunkEnum.TOOL_CALL, block_id=bid, tool_call_id=bid, tool_call_name=tc_name, chunk="") - if btype == "thinking": + if block_type in {"tool_use", "server_tool_use"}: + return cls._chunk( + ChunkEnum.TOOL_CALL, + block_id=bid, + tool_call_id=bid, + tool_call_name=tool_name, + chunk="", + ) + if block_type == "thinking": return cls._chunk(ChunkEnum.THINK, block_id=bid, chunk="") # text or unknown -> CONTENT return cls._chunk(ChunkEnum.CONTENT, block_id=bid, chunk="") @@ -462,46 +316,36 @@ class CcAgentWrapper(BaseAgentWrapper): @classmethod def _message_content_to_chunks( cls, - msg: Any, + msg: "AssistantMessage | UserMessage", session_id: str | None = None, visible_tool_call_ids: set[str] | None = None, include_text: bool = False, ) -> list[StreamChunk]: - """Convert non-partial SDK message content blocks into stream chunks. + """Convert typed SDK content blocks that are not partial events.""" + from claude_agent_sdk import ServerToolResultBlock, TextBlock, ToolResultBlock - Claude Code streams assistant text/tool-use deltas as ``StreamEvent`` - objects, but tool results can arrive later as regular message content - blocks. Surface those blocks so the UI can show what each tool - returned. Some SDK/CLI combinations also put assistant text only in - regular message blocks, so callers can opt into text conversion. - """ chunks: list[StreamChunk] = [] content = getattr(msg, "content", None) - if not isinstance(content, list): + if isinstance(content, str): + if include_text and content: + chunks.append(cls._chunk(ChunkEnum.CONTENT, session_id=session_id, chunk=content)) + return chunks + if content is None: return chunks for block in content: - block_name = block.__class__.__name__ - if include_text and block_name == "TextBlock": - text = getattr(block, "text", "") - if text: - chunks.append(cls._chunk(ChunkEnum.CONTENT, session_id=session_id, chunk=text)) - elif include_text and isinstance(block, dict) and block.get("type") == "text": - text = block.get("text", "") - if text: - chunks.append(cls._chunk(ChunkEnum.CONTENT, session_id=session_id, chunk=text)) - elif include_text and isinstance(block, str): - chunks.append(cls._chunk(ChunkEnum.CONTENT, session_id=session_id, chunk=block)) - elif block_name in {"ToolResultBlock", "ServerToolResultBlock"}: - tool_use_id = getattr(block, "tool_use_id", None) + if include_text and isinstance(block, TextBlock) and block.text: + chunks.append(cls._chunk(ChunkEnum.CONTENT, session_id=session_id, chunk=block.text)) + elif isinstance(block, (ToolResultBlock, ServerToolResultBlock)): + tool_use_id = block.tool_use_id if visible_tool_call_ids is not None and tool_use_id not in visible_tool_call_ids: continue payload: dict[str, Any] = { "tool_use_id": tool_use_id, - "content": getattr(block, "content", None), + "content": block.content, } - if hasattr(block, "is_error"): - payload["is_error"] = getattr(block, "is_error") + if isinstance(block, ToolResultBlock): + payload["is_error"] = block.is_error chunks.append( cls._chunk( ChunkEnum.TOOL_RESULT, @@ -514,25 +358,46 @@ class CcAgentWrapper(BaseAgentWrapper): return chunks - @staticmethod - def _result_message_is_error(msg: Any) -> bool: - """Return whether an SDK ResultMessage represents a failed result.""" - subtype = getattr(msg, "subtype", None) - if isinstance(subtype, str) and subtype.lower() == "success": - return False + @classmethod + def _result_error_text(cls, msg: "ResultMessage") -> str: + """Return the error text used by both the SDK and unified chunks.""" + return "; ".join(msg.errors or []) or str(msg.subtype) - is_error = getattr(msg, "is_error", False) - if isinstance(is_error, bool): - return is_error - if isinstance(is_error, str): - return is_error.lower() in {"true", "error", "errored", "failed", "failure"} - - return isinstance(subtype, str) and subtype.lower() in {"error", "failed", "failure"} - - @staticmethod - def _is_trailing_success_error(exc: Exception) -> bool: - """Return whether an SDK iterator error is the known success-exit artifact.""" - return "Claude Code returned an error result: success" in str(exc) + @classmethod + def _result_message_to_chunks(cls, msg: "ResultMessage") -> list[StreamChunk]: + """Convert the SDK terminal result into usage and error chunks.""" + session_id = msg.session_id or "" + usage = msg.usage or {} + chunks = [ + cls._chunk( + ChunkEnum.USAGE, + session_id=session_id, + chunk=json.dumps(usage), + input_tokens=usage.get("input_tokens"), + output_tokens=usage.get("output_tokens"), + metadata={ + "duration_ms": msg.duration_ms, + "duration_api_ms": msg.duration_api_ms, + "stop_reason": msg.stop_reason, + "num_turns": msg.num_turns, + "total_cost_usd": msg.total_cost_usd, + "model_usage": msg.model_usage, + "permission_denials": msg.permission_denials, + "deferred_tool_use": (asdict(msg.deferred_tool_use) if msg.deferred_tool_use else None), + "api_error_status": msg.api_error_status, + }, + ), + ] + if msg.is_error: + chunks.append( + cls._chunk( + ChunkEnum.ERROR, + session_id=session_id, + chunk=cls._result_error_text(msg), + metadata={"api_error_status": msg.api_error_status}, + ), + ) + return chunks # ----- reply / reply_stream -------------------------------------------- @@ -543,9 +408,10 @@ class CcAgentWrapper(BaseAgentWrapper): opts = self._build_options(inputs, stream=False, **kwargs) last_msg = None - async for msg in query(prompt=inputs, options=opts): - if isinstance(msg, ResultMessage): - last_msg = msg + async with aclosing(query(prompt=inputs, options=opts)) as stream: + async for msg in stream: + if isinstance(msg, ResultMessage): + last_msg = msg if last_msg is None: raise ValueError("No message received from Claude Code.") @@ -561,97 +427,115 @@ class CcAgentWrapper(BaseAgentWrapper): async def reply_stream(self, inputs: Any, **kwargs) -> AsyncGenerator[StreamChunk, None]: """Stream Claude Code events as unified StreamChunk objects.""" - from claude_agent_sdk import query, ResultMessage, AssistantMessage, StreamEvent, UserMessage - from claude_agent_sdk.types import RateLimitEvent + from claude_agent_sdk import ( + AssistantMessage, + MirrorErrorMessage, + query, + RateLimitEvent, + ResultMessage, + StreamEvent, + SystemMessage, + UserMessage, + ) kwargs = self._merged_stream_kwargs(kwargs) opts = self._build_options(inputs, stream=True, **kwargs) - block_ids: dict[int, str] = {} - block_types: dict[int, str] = {} - tool_call_names: dict[int, str] = {} + block_states: dict[int, _BlockState] = {} visible_tool_call_ids: set[str] = set() current_session_id: str | None = None emitted_content = False - received_result_message = False + emitted_reply_end = False + reply_open = False + expected_trailing_error: str | None = None - stream = query(prompt=inputs, options=opts) try: - async for msg in stream: - if isinstance(msg, StreamEvent): - current_session_id = msg.session_id or current_session_id - chunk = self._raw_event_to_chunk( - msg.event, - session_id=msg.session_id, - block_ids=block_ids, - block_types=block_types, - tool_call_names=tool_call_names, - ) - if chunk is not None: - chunk.session_id = chunk.session_id or msg.session_id - if chunk.chunk_type == ChunkEnum.TOOL_CALL and chunk.tool_call_id: - visible_tool_call_ids.add(chunk.tool_call_id) - if chunk.chunk_type == ChunkEnum.CONTENT and chunk.chunk: + async with aclosing(query(prompt=inputs, options=opts)) as stream: + async for msg in stream: + if expected_trailing_error is not None and not ( + isinstance(msg, SystemMessage) and msg.subtype == "session_state_changed" + ): + expected_trailing_error = None + + if isinstance(msg, StreamEvent): + current_session_id = msg.session_id or current_session_id + chunk = self._raw_event_to_chunk( + msg.event, + session_id=msg.session_id, + block_states=block_states, + ) + if chunk is not None: + chunk.session_id = chunk.session_id or msg.session_id + if chunk.chunk_type == ChunkEnum.TOOL_CALL and chunk.tool_call_id: + visible_tool_call_ids.add(chunk.tool_call_id) + if chunk.chunk_type == ChunkEnum.CONTENT and chunk.chunk: + emitted_content = True + if chunk.chunk_type == ChunkEnum.REPLY_START: + reply_open = True + if chunk.chunk_type == ChunkEnum.REPLY_END: + emitted_reply_end = True + reply_open = False + yield chunk + + elif isinstance(msg, UserMessage): + for chunk in self._message_content_to_chunks(msg, current_session_id, visible_tool_call_ids): + yield chunk + + elif isinstance(msg, ResultMessage): + if msg.is_error: + expected_trailing_error = ( + f"Claude Code returned an error result: {self._result_error_text(msg)}" + ) + current_session_id = msg.session_id or current_session_id + if not emitted_content and msg.result: emitted_content = True - yield chunk + yield self._chunk( + ChunkEnum.CONTENT, + session_id=msg.session_id or "", + chunk=msg.result, + ) + for chunk in self._result_message_to_chunks(msg): + yield chunk + if reply_open or not emitted_reply_end: + emitted_reply_end = True + reply_open = False + yield self._chunk( + ChunkEnum.REPLY_END, + session_id=current_session_id, + chunk="", + ) - elif isinstance(msg, UserMessage): - for chunk in self._message_content_to_chunks(msg, current_session_id, visible_tool_call_ids): - yield chunk + elif isinstance(msg, AssistantMessage): + current_session_id = msg.session_id or current_session_id + for chunk in self._message_content_to_chunks( + msg, + current_session_id, + visible_tool_call_ids, + include_text=not emitted_content, + ): + if chunk.chunk_type == ChunkEnum.CONTENT and chunk.chunk: + emitted_content = True + yield chunk - elif isinstance(msg, ResultMessage): - received_result_message = True - current_session_id = msg.session_id or current_session_id - if not emitted_content and getattr(msg, "result", None): - emitted_content = True - yield self._chunk(ChunkEnum.CONTENT, session_id=msg.session_id or "", chunk=msg.result) - # Final result: emit USAGE + REPLY_END - meta = { - "duration_ms": msg.duration_ms, - "duration_api_ms": msg.duration_api_ms, - "stop_reason": msg.stop_reason, - "num_turns": msg.num_turns, - } - yield self._chunk( - ChunkEnum.USAGE, - session_id=msg.session_id or "", - chunk=json.dumps(msg.usage or {}), - metadata=meta, - ) - if self._result_message_is_error(msg): + elif isinstance(msg, MirrorErrorMessage): + self.logger.warning(f"Claude Code session mirror failed: {msg.error}") + yield self._chunk( + ChunkEnum.DATA, + session_id=current_session_id, + chunk=f"Session mirror failed: {msg.error}", + metadata={ + "event": "session_mirror_error", + "session_key": msg.key, + }, + ) + + elif isinstance(msg, RateLimitEvent) and msg.rate_limit_info.status == "rejected": yield self._chunk( ChunkEnum.ERROR, - session_id=msg.session_id or "", - chunk=str(msg.errors) if msg.errors else "Unknown error", + session_id=msg.session_id, + chunk="Rate limit exceeded", ) - yield self._chunk(ChunkEnum.REPLY_END, session_id=msg.session_id or "", chunk="") - - elif isinstance(msg, AssistantMessage): - current_session_id = msg.session_id or current_session_id - # Intermediate assistant text/tool-use is already streamed - # via StreamEvents. Still surface tool-result blocks if the - # SDK includes any in a regular assistant message. - for chunk in self._message_content_to_chunks( - msg, - current_session_id, - visible_tool_call_ids, - include_text=not emitted_content, - ): - if chunk.chunk_type == ChunkEnum.CONTENT and chunk.chunk: - emitted_content = True - yield chunk - - elif isinstance(msg, RateLimitEvent): - yield self._chunk(ChunkEnum.ERROR, session_id=msg.session_id, chunk="Rate limit exceeded") except Exception as exc: - if received_result_message and self._is_trailing_success_error(exc): - self.logger.debug(f"Ignoring Claude Code trailing success error after final result: {exc}") - else: + if expected_trailing_error is None or str(exc) != expected_trailing_error: raise - finally: - try: - await stream.aclose() - except Exception as exc: - if not (received_result_message and self._is_trailing_success_error(exc)): - raise - self.logger.debug(f"Ignoring Claude Code stream close error after final result: {exc}") + self.logger.debug(f"Ignoring Claude Code process exit after error result: {exc}") diff --git a/reme/components/agent_wrapper/cc_session_store.py b/reme/components/agent_wrapper/cc_session_store.py new file mode 100644 index 00000000..ca600c1b --- /dev/null +++ b/reme/components/agent_wrapper/cc_session_store.py @@ -0,0 +1,110 @@ +"""File-backed session store for the Claude Agent SDK.""" + +import json +import shutil +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from claude_agent_sdk import SessionKey, SessionListSubkeysKey, SessionStoreEntry, SessionStoreListEntry + + +class CcFileSessionStore: + """Persist SDK session entries as JSONL under ``project_key/session_id``.""" + + def __init__(self, root: Path) -> None: + self.root = root + + @staticmethod + def _safe_parts(value: str) -> list[str]: + parts = [part for part in value.split("/") if part] + if not parts or any(part in {".", ".."} for part in parts): + raise ValueError(f"Invalid session store path component: {value!r}") + return parts + + def _path(self, *values: str) -> Path: + path = self.root.joinpath(*(part for value in values for part in self._safe_parts(value))) + resolved_root = self.root.resolve() + resolved_path = path.resolve() + if resolved_root != resolved_path and resolved_root not in resolved_path.parents: + raise ValueError(f"Session store path escapes root: {resolved_path}") + return path + + def _project_dir(self, project_key: str) -> Path: + return self._path(project_key) + + def _path_for_key(self, key: "SessionKey") -> Path: + values = [key["project_key"], key["session_id"]] + if subpath := key.get("subpath"): + values.append(subpath) + return self._path(*values).with_suffix(".jsonl") + + @staticmethod + def _read_entries(path: Path) -> list["SessionStoreEntry"]: + entries = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if line.strip(): + entries.append(json.loads(line)) + return entries + + async def append(self, key: "SessionKey", entries: list["SessionStoreEntry"]) -> None: + """Append entries while treating their UUIDs as idempotency keys.""" + if not entries: + return + + path = self._path_for_key(key) + existing = self._read_entries(path) if path.exists() else [] + seen = {entry.get("uuid") for entry in existing if entry.get("uuid")} + new_entries = [] + for entry in entries: + uuid = entry.get("uuid") + if uuid and uuid in seen: + continue + if uuid: + seen.add(uuid) + new_entries.append(entry) + if not new_entries: + return + + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as file: + for entry in new_entries: + file.write(json.dumps(entry, ensure_ascii=False, separators=(",", ":")) + "\n") + + async def load(self, key: "SessionKey") -> list["SessionStoreEntry"] | None: + """Load all entries for a session or subkey.""" + path = self._path_for_key(key) + return self._read_entries(path) if path.exists() else None + + async def list_sessions(self, project_key: str) -> list["SessionStoreListEntry"]: + """List main sessions stored under a project key.""" + project_dir = self._project_dir(project_key) + if not project_dir.is_dir(): + return [] + return [ + {"session_id": path.stem, "mtime": int(path.stat().st_mtime * 1000)} + for path in project_dir.glob("*.jsonl") + if path.is_file() + ] + + async def delete(self, key: "SessionKey") -> None: + """Delete one subkey, or a main session and all of its subkeys.""" + path = self._path_for_key(key) + if path.exists(): + path.unlink() + + if not key.get("subpath"): + session_dir = self._path(key["project_key"], key["session_id"]) + if session_dir.exists(): + shutil.rmtree(session_dir) + + async def list_subkeys(self, key: "SessionListSubkeysKey") -> list[str]: + """List subpaths stored below a main session.""" + session_dir = self._path(key["project_key"], key["session_id"]) + if not session_dir.is_dir(): + return [] + return [ + str(path.relative_to(session_dir).with_suffix("")) + for path in session_dir.rglob("*.jsonl") + if path.is_file() + ] diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index 989830bb..621e62da 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -1,23 +1,42 @@ """Codex Python SDK backend for the unified agent wrapper.""" +from __future__ import annotations + import asyncio from collections.abc import AsyncGenerator from contextlib import suppress -from dataclasses import dataclass, fields, is_dataclass -from enum import Enum +from dataclasses import dataclass +from functools import partial import hashlib import json import os from pathlib import Path import sys import tempfile -from typing import Any +from typing import Any, TYPE_CHECKING from .base_agent_wrapper import BaseAgentWrapper from ..component_registry import R from ...enumeration import ChunkEnum from ...schema import StreamChunk -from ...utils.env_utils import load_env + +if TYPE_CHECKING: + from openai_codex import AsyncCodex, AsyncThread, CodexConfig, RunInput + from openai_codex.types import Notification +else: + # Keep the optional Codex SDK out of ReMe's package import path. Tests may + # also replace this value before the SDK is loaded. + AsyncCodex: Any = None + + +def _get_async_codex_class(): + """Return the optional Codex client class, importing its SDK on first use.""" + global AsyncCodex # pylint: disable=global-statement + if AsyncCodex is None: + from openai_codex import AsyncCodex as AsyncCodexClass + + AsyncCodex = AsyncCodexClass + return AsyncCodex @dataclass(frozen=True) @@ -29,41 +48,69 @@ class _CodexAuthConfig: base_url: str = "" +@dataclass(frozen=True) +class _CodexLaunchConfig: + """Options fixed for the lifetime of one Codex app-server.""" + + auth_mode: str + api_key: str + base_url: str + codex_bin: str | None + config_overrides: tuple[str, ...] + experimental_api: bool + + @R.register("codex") class CodexAgentWrapper(BaseAgentWrapper): """Agent wrapper backed by the Codex Python SDK.""" - def __init__(self, mcp_config: str | None = None, codex_home: str | Path | None = None, **kwargs): + SDK_PACKAGE = "openai-codex" + _CLIENT_OPTION_NAMES = frozenset( + { + "api_key", + "auth_mode", + "base_url", + "codex_bin", + "codex_home", + "config_overrides", + "cwd", + "experimental_api", + "launch_args_override", + }, + ) + + # pylint: disable=too-many-arguments + def __init__( + self, + mcp_config: str | None = None, + codex_home: str | Path | None = None, + *, + auth_mode: str = "auto", + api_key: str = "", + base_url: str = "", + codex_bin: str | None = None, + config_overrides: list[str] | tuple[str, ...] | None = None, + experimental_api: bool = True, + **kwargs, + ) -> None: + if "launch_args_override" in kwargs: + raise TypeError("launch_args_override is not supported; configure codex_bin instead") super().__init__(**kwargs) self.mcp_config = mcp_config self._codex_home = codex_home - self._codex: Any | None = None - self._codex_config: Any | None = None - self._codex_auth_config: _CodexAuthConfig | None = None - self._client_lock = asyncio.Lock() + self._launch_config = _CodexLaunchConfig( + auth_mode=auth_mode, + api_key=api_key, + base_url=base_url, + codex_bin=codex_bin, + config_overrides=tuple(config_overrides or ()), + experimental_api=experimental_api, + ) + self._codex: AsyncCodex | None = None self._turn_lock = asyncio.Lock() self._mcp_snapshot_path: Path | None = None self._thread_tool_contexts: dict[str, str] = {} - @staticmethod - def _first_non_empty(*values: Any) -> str: - for value in values: - if isinstance(value, str) and value: - return value - return "" - - def _default_llm_credential(self) -> dict[str, Any]: - if self.app_context is None: - return {} - from ...enumeration import ComponentEnum - - llm_configs = self.app_context.app_config.components.get(ComponentEnum.AS_LLM) - if not isinstance(llm_configs, dict): - return {} - default_llm = llm_configs.get("default") - credential = getattr(default_llm, "credential", None) - return credential if isinstance(credential, dict) else {} - @property def session_path(self) -> Path: """Directory used for Codex state and persisted threads.""" @@ -153,66 +200,60 @@ class CodexAgentWrapper(BaseAgentWrapper): def _mcp_config_source(self, kwargs: dict[str, Any]) -> str: return self._explicit_mcp_config(kwargs) or str(self._effective_config_snapshot()) - def _resolve_auth_config(self, kwargs: dict[str, Any]) -> _CodexAuthConfig: - """Resolve one explicit Codex auth mode without letting OAuth inherit API credentials.""" - requested_mode = str(kwargs.get("auth_mode") or "auto").lower() + @staticmethod + def _resolve_auth_config(auth_mode: str, api_key: str = "", base_url: str = "") -> _CodexAuthConfig: + """Resolve login from wrapper options. + + The Codex child process still inherits the parent process environment. + """ + requested_mode = str(auth_mode or "auto").lower() if requested_mode not in {"auto", "api_key", "oauth"}: raise ValueError("auth_mode must be one of: auto, api_key, oauth") if requested_mode == "oauth": return _CodexAuthConfig(mode="oauth") - credential = kwargs.get("credential") if isinstance(kwargs.get("credential"), dict) else {} - default_credential = self._default_llm_credential() - api_key = self._first_non_empty( - kwargs.get("api_key"), - credential.get("api_key"), - os.getenv("CODEX_API_KEY"), - os.getenv("OPENAI_API_KEY"), - os.getenv("LLM_API_KEY"), - default_credential.get("api_key"), - ) + api_key = api_key if isinstance(api_key, str) else "" if requested_mode == "api_key" and not api_key: raise ValueError("auth_mode='api_key' requires a non-empty API key") if not api_key: return _CodexAuthConfig(mode="oauth") - base_url = self._first_non_empty( - kwargs.get("base_url"), - credential.get("base_url"), - os.getenv("CODEX_BASE_URL"), - os.getenv("OPENAI_BASE_URL"), - os.getenv("LLM_BASE_URL"), - default_credential.get("base_url"), - ) + base_url = base_url if isinstance(base_url, str) else "" return _CodexAuthConfig(mode="api_key", api_key=api_key, base_url=base_url) - def _build_client_config(self, kwargs: dict[str, Any], auth: _CodexAuthConfig | None = None): + def _build_client_config(self, auth: _CodexAuthConfig) -> CodexConfig: from openai_codex import CodexConfig - auth = auth or self._resolve_auth_config(kwargs) - - project_env = self.project_path / ".env" - env = load_env(project_env) if project_env.exists() else load_env() + env = dict(self.subprocess_environment) self.session_path.mkdir(parents=True, exist_ok=True) env["CODEX_HOME"] = str(self.session_path) - overrides = list(kwargs.get("config_overrides") or []) + overrides = list(self._launch_config.config_overrides) if auth.base_url: overrides.append(f"openai_base_url={json.dumps(auth.base_url)}") - overrides.append(f"forced_login_method={json.dumps('api' if auth.mode == 'api_key' else 'chatgpt')}") + login_method = "api" if auth.mode == "api_key" else "chatgpt" + overrides.append(f"forced_login_method={json.dumps(login_method)}") return CodexConfig( - codex_bin=kwargs.get("codex_bin"), + codex_bin=self._launch_config.codex_bin, config_overrides=tuple(overrides), cwd=str(self.cwd), env=env, client_name="reme", client_title="ReMe", + experimental_api=self._launch_config.experimental_api, ) + @classmethod + def _reject_client_options(cls, kwargs: dict[str, Any]) -> None: + invalid = sorted(cls._CLIENT_OPTION_NAMES.intersection(kwargs)) + if invalid: + names = ", ".join(invalid) + raise TypeError(f"Codex client options must be configured on the wrapper: {names}") + def _mcp_server_config(self, kwargs: dict[str, Any]) -> dict[str, Any] | None: from ..job import BackgroundJob, StreamJob - job_names = list(kwargs.get("job_tools") or []) + job_names = list(dict.fromkeys(kwargs.get("job_tools") or [])) if not job_names: return None jobs = self._resolve_job_tools(job_names) @@ -221,20 +262,20 @@ class CodexAgentWrapper(BaseAgentWrapper): raise TypeError(f"Codex job_tools must be non-stream request jobs: {', '.join(unsupported)}") config_source = self._mcp_config_source(kwargs) + args = [ + "-m", + "reme.components.agent_wrapper.codex_mcp_server", + "--config", + config_source, + "--workspace", + str(self.workspace_path), + ] + for name in job_names: + args.extend(["--job", name]) + args.extend(["--tool-context-id", str(kwargs.get("tool_context_id") or "")]) return { "command": sys.executable, - "args": [ - "-m", - "reme.components.agent_wrapper.codex_mcp_server", - "--config", - config_source, - "--workspace", - str(self.workspace_path), - "--jobs", - json.dumps(job_names), - "--tool-context-id", - str(kwargs.get("tool_context_id") or ""), - ], + "args": args, "cwd": str(self.project_path), "required": True, "enabled_tools": job_names, @@ -257,8 +298,9 @@ class CodexAgentWrapper(BaseAgentWrapper): return default return value if isinstance(value, enum_cls) else enum_cls(value) - async def _open_thread(self, codex: Any, kwargs: dict[str, Any]): + async def _open_thread(self, codex: AsyncCodex, kwargs: dict[str, Any]) -> AsyncThread: from openai_codex import ApprovalMode, Sandbox + from openai_codex.types import Personality, ThreadSource, ThreadStartSource resume = kwargs.get("resume") or "" session_id = kwargs.get("session_id") or "" @@ -284,18 +326,32 @@ class CodexAgentWrapper(BaseAgentWrapper): "sandbox": self._enum(Sandbox, kwargs.get("sandbox"), Sandbox.full_access), "service_tier": kwargs.get("service_tier"), } + personality = self._enum(Personality, kwargs.get("personality")) + thread_source = self._enum(ThreadSource, kwargs.get("thread_source")) if fork_session: - thread = await codex.thread_fork(thread_id, ephemeral=kwargs.get("ephemeral"), **common) + thread = await codex.thread_fork( + thread_id, + ephemeral=kwargs.get("ephemeral"), + thread_source=thread_source, + **common, + ) elif thread_id: - thread = await codex.thread_resume(thread_id, **common) + thread = await codex.thread_resume(thread_id, personality=personality, **common) else: - thread = await codex.thread_start(ephemeral=kwargs.get("ephemeral"), **common) + thread = await codex.thread_start( + ephemeral=kwargs.get("ephemeral"), + personality=personality, + service_name=kwargs.get("service_name"), + session_start_source=self._enum(ThreadStartSource, kwargs.get("session_start_source")), + thread_source=thread_source, + **common, + ) self._thread_tool_contexts[thread.id] = requested_tool_context return thread def _turn_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: from openai_codex import ApprovalMode, Sandbox - from openai_codex.generated.v2_all import Personality, ReasoningEffort, ReasoningSummary + from openai_codex.types import Personality, ReasoningEffort, ReasoningSummary return { "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode")), @@ -309,39 +365,33 @@ class CodexAgentWrapper(BaseAgentWrapper): "summary": self._enum(ReasoningSummary, kwargs.get("summary")), } - async def _get_codex(self, kwargs: dict[str, Any]) -> Any: - """Lazily start one app-server and reject launch-config changes while it is live.""" - from openai_codex import AsyncCodex - - auth = self._resolve_auth_config(kwargs) - config = self._build_client_config(kwargs, auth) - async with self._client_lock: - if self._codex is not None: - if config != self._codex_config or auth != self._codex_auth_config: - raise RuntimeError("Codex client configuration changed; close the wrapper before reconfiguring it") - return self._codex - codex = AsyncCodex(config) - try: - if auth.mode == "api_key": - await codex.login_api_key(auth.api_key) - else: - account = await codex.account() - if account.account is None: - raise RuntimeError(f"No ChatGPT OAuth login found in CODEX_HOME: {self.session_path}") - except BaseException: - await codex.close() - raise - self._codex = codex - self._codex_config = config - self._codex_auth_config = auth - return codex + async def _get_codex(self) -> AsyncCodex: + """Lazily start the component-owned client from its fixed launch configuration.""" + if self._codex is not None: + return self._codex + auth = self._resolve_auth_config( + self._launch_config.auth_mode, + self._launch_config.api_key, + self._launch_config.base_url, + ) + codex = _get_async_codex_class()(self._build_client_config(auth)) + try: + if auth.mode == "api_key": + await codex.login_api_key(auth.api_key) + else: + account = await codex.account() + if account.account is None: + raise RuntimeError(f"No ChatGPT OAuth login found in CODEX_HOME: {self.session_path}") + except BaseException: + await codex.close() + raise + self._codex = codex + return codex async def _close(self) -> None: """Close the persistent app-server and remove its private config snapshot.""" - async with self._client_lock: + async with self._turn_lock: codex, self._codex = self._codex, None - self._codex_config = None - self._codex_auth_config = None self._thread_tool_contexts.clear() try: if codex is not None: @@ -351,28 +401,20 @@ class CodexAgentWrapper(BaseAgentWrapper): self._mcp_snapshot_path.unlink(missing_ok=True) self._mcp_snapshot_path = None - @classmethod - def _serialize(cls, value: Any) -> Any: - if hasattr(value, "model_dump"): - return value.model_dump(mode="json", by_alias=True) - if isinstance(value, Enum): - return value.value - if is_dataclass(value): - return {field.name: cls._serialize(getattr(value, field.name)) for field in fields(value)} - if isinstance(value, dict): - return {key: cls._serialize(item) for key, item in value.items()} - if isinstance(value, (list, tuple)): - return [cls._serialize(item) for item in value] - return value + @staticmethod + def _serialize(value: Any) -> Any: + from pydantic_core import to_jsonable_python - async def reply(self, inputs: Any, **kwargs) -> dict: + return to_jsonable_python(value, by_alias=True) + + async def reply(self, inputs: RunInput, **kwargs) -> dict: """Run one Codex turn and return its final response.""" - if not isinstance(inputs, str): - raise NotImplementedError("Only string input is supported for Codex.") + self._reject_client_options(kwargs) kwargs = self._merged_kwargs(kwargs) self._ensure_skills(kwargs.get("skills")) + await self.start() async with self._turn_lock: - codex = await self._get_codex(kwargs) + codex = await self._get_codex() thread = await self._open_thread(codex, kwargs) result = await thread.run(inputs, **self._turn_kwargs(kwargs)) @@ -390,27 +432,22 @@ class CodexAgentWrapper(BaseAgentWrapper): raise ValueError("Codex returned invalid JSON for the requested output_schema") from exc return response - @classmethod - def _item_data(cls, item: Any) -> tuple[Any, str, str]: - item = item.root if hasattr(item, "root") else item - return item, getattr(item, "type", ""), getattr(item, "id", "") - @classmethod # pylint: disable=too-many-return-statements - def _event_to_chunks(cls, event: Any, session_id: str) -> list[StreamChunk]: + def _event_to_chunks(cls, event: Notification, session_id: str) -> list[StreamChunk]: """Convert one Codex app-server notification to unified stream chunks.""" method, payload = event.method, event.payload + make_chunk = partial(cls._chunk, session_id=session_id) if method == "turn/started": - return [cls._chunk(ChunkEnum.REPLY_START, session_id=session_id, metadata={"turn_id": payload.turn.id})] + return [make_chunk(ChunkEnum.REPLY_START, metadata={"turn_id": payload.turn.id})] if method == "item/agentMessage/delta": - return [cls._chunk(ChunkEnum.CONTENT, session_id=session_id, block_id=payload.item_id, chunk=payload.delta)] + return [make_chunk(ChunkEnum.CONTENT, block_id=payload.item_id, chunk=payload.delta)] if method in {"item/reasoning/summaryTextDelta", "item/reasoning/textDelta", "item/plan/delta"}: - return [cls._chunk(ChunkEnum.THINK, session_id=session_id, block_id=payload.item_id, chunk=payload.delta)] + return [make_chunk(ChunkEnum.THINK, block_id=payload.item_id, chunk=payload.delta)] if method in {"item/commandExecution/outputDelta", "item/fileChange/outputDelta"}: return [ - cls._chunk( + make_chunk( ChunkEnum.TOOL_RESULT, - session_id=session_id, block_id=payload.item_id, tool_call_id=payload.item_id, chunk=payload.delta, @@ -418,9 +455,8 @@ class CodexAgentWrapper(BaseAgentWrapper): ] if method == "item/mcpToolCall/progress": return [ - cls._chunk( + make_chunk( ChunkEnum.TOOL_RESULT, - session_id=session_id, block_id=payload.item_id, tool_call_id=payload.item_id, chunk=payload.message, @@ -434,9 +470,8 @@ class CodexAgentWrapper(BaseAgentWrapper): status = "started" if method.endswith("/started") else "completed" decision_source = cls._serialize(getattr(payload, "decision_source", None)) return [ - cls._chunk( + make_chunk( ChunkEnum.APPROVAL, - session_id=session_id, block_id=target_item_id or review_id, tool_call_id=target_item_id, chunk=action, @@ -450,16 +485,16 @@ class CodexAgentWrapper(BaseAgentWrapper): ), ] if method in {"item/started", "item/completed"}: - item, item_type, item_id = cls._item_data(payload.item) + item = payload.item.root + item_type, item_id = item.type, item.id tool_types = {"commandExecution", "fileChange", "mcpToolCall", "dynamicToolCall", "collabAgentToolCall"} if item_type not in tool_types: return [] name = getattr(item, "tool", None) or item_type chunk_type = ChunkEnum.TOOL_CALL if method == "item/started" else ChunkEnum.TOOL_RESULT return [ - cls._chunk( + make_chunk( chunk_type, - session_id=session_id, block_id=item_id, tool_call_id=item_id, tool_call_name=name, @@ -470,9 +505,8 @@ class CodexAgentWrapper(BaseAgentWrapper): usage = payload.token_usage.last data = cls._serialize(usage) return [ - cls._chunk( + make_chunk( ChunkEnum.USAGE, - session_id=session_id, chunk=data, input_tokens=usage.input_tokens, output_tokens=usage.output_tokens, @@ -480,9 +514,8 @@ class CodexAgentWrapper(BaseAgentWrapper): ] if method == "error": return [ - cls._chunk( + make_chunk( ChunkEnum.ERROR, - session_id=session_id, chunk=payload.error.message, metadata={"will_retry": payload.will_retry}, ), @@ -490,30 +523,38 @@ class CodexAgentWrapper(BaseAgentWrapper): if method == "turn/completed": turn = payload.turn chunks = [] - if getattr(turn, "error", None): - chunks.append(cls._chunk(ChunkEnum.ERROR, session_id=session_id, chunk=turn.error.message)) + if turn.error: + chunks.append(make_chunk(ChunkEnum.ERROR, chunk=turn.error.message)) chunks.append( - cls._chunk( + make_chunk( ChunkEnum.REPLY_END, - session_id=session_id, metadata={ "turn_id": turn.id, - "status": getattr(turn.status, "value", str(turn.status)), + "status": turn.status.value, "duration_ms": turn.duration_ms, }, ), ) return chunks + if getattr(payload, "turn_id", None): + return [ + make_chunk( + ChunkEnum.DATA, + block_id=getattr(payload, "item_id", None), + chunk=cls._serialize(payload), + metadata={"codex_method": method}, + ), + ] return [] - async def reply_stream(self, inputs: Any, **kwargs) -> AsyncGenerator[StreamChunk, None]: + async def reply_stream(self, inputs: RunInput, **kwargs) -> AsyncGenerator[StreamChunk, None]: """Stream Codex app-server notifications as unified chunks.""" - if not isinstance(inputs, str): - raise NotImplementedError("Only string input is supported for Codex.") + self._reject_client_options(kwargs) kwargs = self._merged_stream_kwargs(kwargs) self._ensure_skills(kwargs.get("skills")) + await self.start() async with self._turn_lock: - codex = await self._get_codex(kwargs) + codex = await self._get_codex() thread = await self._open_thread(codex, kwargs) turn = await thread.turn(inputs, **self._turn_kwargs(kwargs)) stream = turn.stream() diff --git a/reme/components/agent_wrapper/codex_mcp_server.py b/reme/components/agent_wrapper/codex_mcp_server.py index 16e43704..7ddfbe59 100644 --- a/reme/components/agent_wrapper/codex_mcp_server.py +++ b/reme/components/agent_wrapper/codex_mcp_server.py @@ -1,15 +1,9 @@ """FastMCP STDIO bridge that exposes selected ReMe jobs to Codex.""" import argparse -from contextlib import asynccontextmanager -import json from pathlib import Path from typing import Any -from fastmcp import FastMCP -from fastmcp.exceptions import ToolError -from fastmcp.tools import FunctionTool - from ...config import resolve_app_config from ...reme import ReMe @@ -18,61 +12,46 @@ def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Expose selected ReMe jobs over FastMCP STDIO") parser.add_argument("--config", default="default", help="ReMe config name or file path") parser.add_argument("--workspace", required=True, help="ReMe workspace directory") - parser.add_argument("--jobs", required=True, help="JSON array of ReMe job names") + parser.add_argument("--job", dest="jobs", action="append", required=True, help="ReMe job name; repeat as needed") parser.add_argument("--tool-context-id", default="", help="Context id injected into every job call") return parser.parse_args() -def _load_job_names(raw: str) -> list[str]: - value = json.loads(raw) - if not isinstance(value, list) or not all(isinstance(name, str) and name for name in value): - raise ValueError("--jobs must be a JSON array of non-empty strings") - return value +def _prepare_config(config: dict[str, Any], job_names: list[str], tool_context_id: str = "") -> dict[str, Any]: + """Configure the dedicated child Application to serve selected jobs over MCP STDIO.""" + selected = set(job_names) + jobs: dict[str, dict[str, Any]] = {} + for name, raw_job_config in (config.get("jobs") or {}).items(): + job_config = dict(raw_job_config) + if job_config.get("backend") in {"background", "cron"}: + continue + if name in selected: + # An explicit Codex job_tools selection has always overridden enable_serve. + job_config["enable_serve"] = True + jobs[name] = job_config + missing = sorted(selected.difference(jobs)) + if missing: + raise KeyError(f"Codex job tools not found or not request jobs: {', '.join(missing)}") -def _make_tool(job: Any, tool_context_id: str) -> FunctionTool: - async def execute_tool(**kwargs): - if tool_context_id: - if "tool_context_id" in kwargs: - raise ToolError("tool_context_id is managed by the Codex agent wrapper") - kwargs["tool_context_id"] = tool_context_id - response = await job(**kwargs) - if not response.success: - raise ToolError(str(response.answer)) - return response.answer + service: dict[str, Any] = { + "backend": "mcp", + "transport": "stdio", + "jobs": job_names, + "tool_error_on_failure": True, + } + if tool_context_id: + service["injected_job_kwargs"] = {"tool_context_id": tool_context_id} - return FunctionTool( - name=job.name, - description=job.description, - fn=execute_tool, - parameters=job.parameters or {}, - ) - - -def build_server(app: ReMe, job_names: list[str], tool_context_id: str = "") -> FastMCP: - """Build a STDIO server backed by a dedicated ReMe Application.""" - - @asynccontextmanager - async def lifespan(_server): - await app.start() - try: - yield - finally: - await app.close() - - server = FastMCP(name="reme-codex-tools", lifespan=lifespan) - for name in job_names: - job = app.context.jobs.get(name) - if job is None: - raise KeyError(f"Job '{name}' not found") - server.add_tool(_make_tool(job, tool_context_id)) - return server + prepared = dict(config) + prepared["jobs"] = jobs + prepared["service"] = service + return prepared def main() -> None: """Load ReMe and serve the requested jobs over STDIO.""" args = _parse_args() - job_names = _load_job_names(args.jobs) config = resolve_app_config( config=args.config, workspace_dir=str(Path(args.workspace).absolute()), @@ -81,16 +60,8 @@ def main() -> None: log_to_file=False, log_config=False, ) - # The bridge needs ordinary jobs available for nested job references, but - # must not start workspace watchers or cron loops in this short-lived child. - config["jobs"] = { - name: job_config - for name, job_config in (config.get("jobs") or {}).items() - if job_config.get("backend") not in {"background", "cron"} - } - app = ReMe(**config) - server = build_server(app, job_names, args.tool_context_id) - server.run(transport="stdio", show_banner=False) + config = _prepare_config(config, args.jobs, args.tool_context_id) + ReMe(**config).run_app() if __name__ == "__main__": diff --git a/reme/components/application_context.py b/reme/components/application_context.py index 611c188a..317e81e7 100644 --- a/reme/components/application_context.py +++ b/reme/components/application_context.py @@ -29,9 +29,9 @@ class ApplicationContext: self.components: dict[ComponentEnum, dict[str, "BaseComponent"]] = {} self.jobs: dict[str, "BaseJob"] = {} self.thread_pool: ThreadPoolExecutor | None = None - # Side-channel for service/transport-specific objects that don't fit - # the shared component/job model — e.g. MCPService publishes a - # ChannelSink under "channel_sink" so MCP-specific steps - # (claim_channel, channel_notify) can find it. Keep keys narrow: - # if a value is needed across services, promote it to a typed field. + + # Application-lifetime shared state. Values remain available across Job and Step + # invocations while this Application is running, so Jobs may keep cross-call state here. + # This is in-memory state, not durable storage; use workspace files or a store when state + # must survive an Application restart. self.metadata: dict[str, Any] = {} diff --git a/reme/components/service/base_service.py b/reme/components/service/base_service.py index d34745b8..c1afd2cf 100644 --- a/reme/components/service/base_service.py +++ b/reme/components/service/base_service.py @@ -20,8 +20,9 @@ class BaseService(BaseComponent): component_type = ComponentEnum.SERVICE - def __init__(self, **kwargs): + def __init__(self, jobs: list[str] | None = None, **kwargs): super().__init__(**kwargs) + self.jobs: set[str] | None = set(jobs) if jobs is not None else None # Underlying framework instance (FastAPI, FastMCP, ...); populated by build_service(). self.service = None @@ -55,26 +56,42 @@ class BaseService(BaseComponent): @asynccontextmanager async def lifespan(_): await app.start() - service_info = json.dumps({"host": host, "port": port}) - os.environ[REME_SERVICE_INFO] = service_info - self.logger.info(f"{self.name} started: {REME_SERVICE_INFO}={service_info}") - yield - await app.close() + try: + service_info = json.dumps({"host": host, "port": port}) + os.environ[REME_SERVICE_INFO] = service_info + self.logger.info(f"{self.name} started: {REME_SERVICE_INFO}={service_info}") + yield + finally: + await app.close() return lifespan def add_jobs(self, app: "Application") -> None: - """Register every job whose enable_serve flag is True.""" + """Register service-enabled jobs, optionally restricted by the configured whitelist.""" + if self.jobs is not None: + missing = sorted(self.jobs.difference(app.context.jobs)) + if missing: + raise KeyError(f"Service jobs not found: {', '.join(missing)}") + disabled = sorted(name for name in self.jobs if not app.context.jobs[name].enable_serve) + if disabled: + raise ValueError(f"Service jobs are not enabled for serving: {', '.join(disabled)}") + for name, job in app.context.jobs.items(): - if not job.enable_serve: + if not job.enable_serve or (self.jobs is not None and name not in self.jobs): continue try: - if self.add_job(job): - self.logger.info(f"Added job: {name}") - else: - self.logger.warning(f"Skipped job: {name}") + added = self.add_job(job) except Exception as e: + if self.jobs is not None: + raise self.logger.error(f"Failed to add job {name}: {e}") + continue + if added: + self.logger.info(f"Added job: {name}") + elif self.jobs is not None: + raise TypeError(f"Service does not support job: {name}") + else: + self.logger.warning(f"Skipped job: {name}") def run_app(self, app: "Application") -> None: """Build the service, register jobs, then start serving (blocking).""" diff --git a/reme/components/service/mcp_service.py b/reme/components/service/mcp_service.py index a8fc3558..4de5b94d 100644 --- a/reme/components/service/mcp_service.py +++ b/reme/components/service/mcp_service.py @@ -1,112 +1,17 @@ -"""MCP service: expose jobs as MCP tools. +"""MCP service: expose jobs as MCP tools.""" -Channel binding (the `` -push from background steps to a specific Claude Code window) is uniform -across transports: a single `ChannelSink` lives on -`ApplicationContext.metadata["channel_sink"]`, unbound at startup, and any -client calling the `claim_channel` MCP tool binds itself as the recipient -via `fastmcp.server.dependencies.get_context().session`. Last-claim-wins. - -Under stdio (one client per server process) the client should claim once -after init; until then channel events drop silently. Under shared -streamable-http / sse the human picks which window receives events. - -``ChannelSink`` is colocated here because it is the runtime mechanism -behind this service's channel feature — pushes ``notifications/claude/channel`` -frames to the bound MCP session. Lossy by design: not bound → no-op; -``send_message`` raises → log warning, swallow (failed notifications must -not surface as ingest failures). Uses ``ServerSession.send_message`` -(low-level raw frame) instead of ``send_notification`` because the latter -validates against a closed ``ServerNotification`` RootModel union that -does not include ``notifications/claude/channel`` — Pydantic rejects -custom methods. Meta keys are filtered to ``[A-Za-z0-9_]+``: Claude Code -silently drops keys with hyphens / other chars when projecting onto -```` attrs. -""" - -import re -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from .base_service import BaseService from ..component_registry import R from ..job import BaseJob, StreamJob from ...constants import REME_DEFAULT_HOST, REME_DEFAULT_PORT -from ...utils import get_logger if TYPE_CHECKING: from fastmcp.server.server import Transport - from mcp.server.session import ServerSession - from ...application import Application -_IDENT_RE = re.compile(r"^[A-Za-z0-9_]+$") -_CHANNEL_METHOD = "notifications/claude/channel" - - -class ChannelSink: - """Hold a bound MCP ``ServerSession`` and forward channel events to it.""" - - def __init__(self) -> None: - self._session: "ServerSession | None" = None - self._logger = get_logger() - - def bind(self, session: "ServerSession") -> None: - """Set ``session`` as the recipient for subsequent ``emit`` calls (last-claim-wins).""" - self._session = session - - def unbind(self) -> None: - """Drop the bound session; future ``emit`` calls become no-ops until rebind.""" - self._session = None - - async def emit(self, content: str, meta: dict[str, str] | None = None) -> None: - """Send one channel notification; no-op if unbound, log+swallow on transport failure.""" - session = self._session - if session is None: - return - - from mcp.shared.message import SessionMessage - from mcp.types import JSONRPCMessage, JSONRPCNotification - - clean_meta = {k: str(v) for k, v in (meta or {}).items() if _IDENT_RE.match(k)} - message = SessionMessage( - JSONRPCMessage( - JSONRPCNotification( - jsonrpc="2.0", - method=_CHANNEL_METHOD, - params={"content": content, "meta": clean_meta}, - ), - ), - ) - - try: - await session.send_message(message) - except Exception as exc: - self._logger.warning(f"ChannelSink: send_message failed ({type(exc).__name__}: {exc})") - - -_CHANNEL_INSTRUCTIONS = ( - "Events from the reme channel arrive as\n" - ' \n' - " added|modified|deleted: \n" - " ...\n" - " \n" - "The workspace watcher fires one event per debounced batch of live changes " - "under daily/, digest/, and resource/ (initial-scan diffs at startup are " - "intentionally NOT replayed).\n" - "\n" - "Events are delivered ONLY to the MCP session that called the " - "`claim_channel` tool last. Call it once per Claude Code window that " - "should receive workspace-change notifications.\n" - "\n" - "When new files appear under daily/ or resource/, treat it as a suggestion " - "to run `/dream ` on each new path -- unless the user is mid-task and " - "would be interrupted, in which case acknowledge in one line and continue. " - "For changes under digest/ (which /dream itself writes), just acknowledge; " - "do not re-dream them. For deletes or modifies elsewhere, just acknowledge." -) - - @R.register("mcp") class MCPService(BaseService): """Expose non-stream jobs as MCP tools over stdio, SSE, or streamable-http.""" @@ -116,43 +21,62 @@ class MCPService(BaseService): transport: "Transport" = "sse", host: str = REME_DEFAULT_HOST, port: int = REME_DEFAULT_PORT, + injected_job_kwargs: dict[str, Any] | None = None, + tool_error_on_failure: bool = False, **kwargs, ): super().__init__(**kwargs) self.transport: Transport = transport self.host: str = host self.port: int = port + self.injected_job_kwargs = dict(injected_job_kwargs or {}) + self.tool_error_on_failure = tool_error_on_failure # ----- BaseService contract ------------------------------------------ def build_service(self, app: "Application") -> None: - """Construct the FastMCP server and publish an unbound ChannelSink.""" + """Construct the FastMCP server.""" from fastmcp import FastMCP - app.context.metadata["channel_sink"] = ChannelSink() self.service = FastMCP( name=app.config.app_name, - instructions=_CHANNEL_INSTRUCTIONS, lifespan=self._lifespan(app, self.host, self.port), ) def add_job(self, job: BaseJob) -> bool: """Register a non-stream job as an MCP tool; StreamJobs are unsupported.""" + from fastmcp.exceptions import ToolError from fastmcp.tools import FunctionTool if isinstance(job, StreamJob): return False async def execute_tool(**kwargs): + conflicts = sorted(self.injected_job_kwargs.keys() & kwargs.keys()) + if conflicts: + names = ", ".join(conflicts) + raise ToolError(f"{names} injected by the MCP server and cannot be provided by the caller") + kwargs.update(self.injected_job_kwargs) response = await job(**kwargs) + if self.tool_error_on_failure and not response.success: + raise ToolError(str(response.answer)) return response.answer + parameters = dict(job.parameters or {}) + injected_names = self.injected_job_kwargs.keys() + if "properties" in parameters: + parameters["properties"] = { + name: schema for name, schema in parameters["properties"].items() if name not in injected_names + } + if "required" in parameters: + parameters["required"] = [name for name in parameters["required"] if name not in injected_names] + self.service.add_tool( FunctionTool( name=job.name, description=job.description, fn=execute_tool, - parameters=job.parameters or {}, + parameters=parameters, ), ) return True diff --git a/reme/config/default.yaml b/reme/config/default.yaml index d541aa8e..18061b6b 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -672,7 +672,6 @@ components: api_key: ${CLAUDE_CODE_API_KEY:-} base_url: ${CLAUDE_CODE_BASE_URL:-https://dashscope.aliyuncs.com/apps/anthropic} permission_mode: bypassPermissions - system_prompt_mode: replace codex: backend: codex auth_mode: api_key diff --git a/reme/config/demo.yaml b/reme/config/demo.yaml index 1f6b2ae9..b883ce83 100644 --- a/reme/config/demo.yaml +++ b/reme/config/demo.yaml @@ -39,6 +39,44 @@ jobs: steps: - backend: add_step + llm_demo: + backend: base + description: "LLM agent demo with the add job available as a tool" + query: "Use the add tool to calculate 1 + 2 and return the result." + sys_prompt: "Use the add tool for arithmetic tasks." + parameters: + type: object + properties: + query: + type: string + description: "message sent to the agent" + default: "Use the add tool to calculate 1 + 2 and return the result." + sys_prompt: + type: string + description: "optional system prompt for the agent" + steps: + - backend: llm_demo_step + agent_wrapper: codex + + stream_llm_demo: + backend: stream + description: "streaming LLM agent demo with the add job available as a tool" + query: "Use the add tool to calculate 1 + 2 and return the result." + sys_prompt: "Use the add tool for arithmetic tasks." + parameters: + type: object + properties: + query: + type: string + description: "message sent to the agent" + default: "Use the add tool to calculate 1 + 2 and return the result." + sys_prompt: + type: string + description: "optional system prompt for the agent" + steps: + - backend: stream_llm_demo_step + agent_wrapper: codex + stream_demo: backend: stream description: "stream demo job: repeat query 10x and stream char-by-char" @@ -61,3 +99,51 @@ jobs: steps: - backend: stream_demo_step1 - backend: stream_demo_step2 + +components: + as_llm: + default: + backend: openai + model: qwen3.7-max + stream: true + context_size: 200000 + max_retries: 3 + credential: + api_key: ${LLM_API_KEY:-} + base_url: ${LLM_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} + parameters: + max_tokens: 65536 + thinking_enable: false + + agent_wrapper: + as: + backend: agentscope + as_llm: default + job_tools: + - add + permission_mode: bypass + react_config: + max_iters: 100 + context_config: + trigger_ratio: 0.89 + reserve_ratio: 0.1 + tool_result_limit: 50000 + model_config: + max_retries: 3 + cc: + backend: claude_code + model: ${CLAUDE_CODE_MODEL_NAME:-qwen3.7-max} + api_key: ${CLAUDE_CODE_API_KEY:-} + base_url: ${CLAUDE_CODE_BASE_URL:-https://dashscope.aliyuncs.com/apps/anthropic} + job_tools: + - add + permission_mode: bypassPermissions + codex: + backend: codex + auth_mode: oauth + model: ${CODEX_MODEL_NAME:-} + codex_home: ${CODEX_HOME:-~/.codex} + job_tools: + - add + approval_mode: auto_review + sandbox: full-access diff --git a/reme/enumeration/chunk_enum.py b/reme/enumeration/chunk_enum.py index 4e7d6499..eb855371 100644 --- a/reme/enumeration/chunk_enum.py +++ b/reme/enumeration/chunk_enum.py @@ -21,12 +21,13 @@ class ChunkEnum(str, Enum): Claude Code SDK events -> ChunkEnum mapping: message_start -> REPLY_START + message_delta -> USAGE message_stop -> REPLY_END content_block_start/delta/stop (text) -> CONTENT content_block_start/delta/stop (thinking) -> THINK content_block_start/delta/stop (tool_use) -> TOOL_CALL ToolResultBlock -> TOOL_RESULT - ResultMessage -> USAGE + DONE + ResultMessage -> USAGE ResultMessage.is_error -> ERROR Codex app-server notifications follow the same lifecycle categories; diff --git a/reme/reme.py b/reme/reme.py index 6af9da89..a1f98efe 100644 --- a/reme/reme.py +++ b/reme/reme.py @@ -61,10 +61,11 @@ async def call_server(action: str, **kwargs): def main(): """Parse CLI arguments and launch the appropriate mode.""" - load_env() + environment = load_env() action, kwargs = parse_args(*sys.argv[1:]) if action == "start": kwargs = prepare_start_config(kwargs) + kwargs["environment"] = environment if should_precheck_start(kwargs) and not precheck_start(kwargs.get("service")): return ReMe(**kwargs).run_app() diff --git a/reme/schema/application_config.py b/reme/schema/application_config.py index a3a8c8bf..2454e6b9 100644 --- a/reme/schema/application_config.py +++ b/reme/schema/application_config.py @@ -28,6 +28,10 @@ class ApplicationConfig(BaseModel): """Root config for the ReMe application.""" app_name: str = Field(default=os.getenv("APP_NAME", "ReMe"), description="Application display name") + environment: dict[str, str] = Field( + default_factory=dict, + description="Environment variables loaded once at startup and passed to agent subprocesses", + ) workspace_dir: str = Field(default=".reme", description="Workspace root directory for runtime files") metadata_dir: str = Field(default="metadata", description="Subdirectory for ReMe persistent state") session_dir: str = Field(default="session", description="Subdirectory for persisted agent sessions") diff --git a/reme/steps/__init__.py b/reme/steps/__init__.py index 2cd2b69e..80d1d8c0 100644 --- a/reme/steps/__init__.py +++ b/reme/steps/__init__.py @@ -1,12 +1,11 @@ """steps""" -from . import benchmark, channel, common, evolve, file_io, index, transfer +from . import benchmark, common, evolve, file_io, index, transfer from .base_step import BaseStep __all__ = [ "BaseStep", "benchmark", - "channel", "common", "evolve", "file_io", diff --git a/reme/steps/channel/__init__.py b/reme/steps/channel/__init__.py deleted file mode 100644 index 51062e80..00000000 --- a/reme/steps/channel/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -"""Channel steps.""" - -from .channel_notify import ChannelNotifyStep -from .claim_channel import ClaimChannelStep - -__all__ = [ - "ChannelNotifyStep", - "ClaimChannelStep", -] diff --git a/reme/steps/channel/channel_notify.py b/reme/steps/channel/channel_notify.py deleted file mode 100644 index e19f877d..00000000 --- a/reme/steps/channel/channel_notify.py +++ /dev/null @@ -1,68 +0,0 @@ -"""``channel_notify_step`` — push a debounced batch of workspace changes as one channel event. - -Designed to be slotted into ``watch_changes_step.dispatch_steps`` next to -``update_index_step``: when the watcher emits a batch of changes, this -step forwards a single human-readable summary to -``ApplicationContext.metadata["channel_sink"]``. The Claude Code main -session then sees a ```` -tag and reacts per the server's ``instructions``. - -One event per batch (not per file) — the watcher already de-bounces and -de-duplicates, so a batch is a meaningful "things that changed together" -unit. Putting N events on the wire per batch would multiply session -turns without adding signal. - -No-op (silently) when: - -* ``channel_sink`` is absent from the application context metadata - (e.g. service wasn't an ``MCPService``), so this step is safe in - any pipeline. -* ``context['changes']`` is missing or empty. -* The sink itself has no bound session (no client called ``claim_channel``). -""" - -from pathlib import Path - -from ..base_step import BaseStep -from ...components import R - - -@R.register("channel_notify_step") -class ChannelNotifyStep(BaseStep): - """Forward a batch of workspace changes to the Claude Code channel.""" - - async def execute(self): - sink = self.app_context.metadata.get("channel_sink") if self.app_context is not None else None - if sink is None: - return self.context.response if self.context is not None else None - - changes = (self.context.get("changes", []) if self.context is not None else []) or [] - if not changes: - return self.context.response if self.context is not None else None - - # Render paths workspace-relative so the agent can pass them directly to - # slash commands like /dream . Absolute paths that fall outside - # the workspace are left as-is rather than erroring — they shouldn't occur, - # but a stray entry shouldn't kill the event. - workspace = self.workspace_path - lines: list[str] = [] - for change in changes: - try: - raw = Path(change["path"]) - except (KeyError, TypeError): - continue - try: - shown = str(raw.resolve().relative_to(workspace)) - except ValueError: - shown = str(raw) - lines.append(f"{change.get('change', '?')}: {shown}") - - if not lines: - return self.context.response if self.context is not None else None - - self.logger.info(f"[channel_notify] emit batch count={len(lines)}") - await sink.emit( - content="Workspace 变更:\n" + "\n".join(lines), - meta={"kind": "workspace_change", "count": str(len(lines))}, - ) - return self.context.response if self.context is not None else None diff --git a/reme/steps/channel/claim_channel.py b/reme/steps/channel/claim_channel.py deleted file mode 100644 index 65edbb5d..00000000 --- a/reme/steps/channel/claim_channel.py +++ /dev/null @@ -1,60 +0,0 @@ -"""``claim_channel_step`` — let an MCP client elect itself as the ```` recipient. - -The single bind path for every transport (stdio, sse, streamable-http): -this step uses ``fastmcp.server.dependencies.get_context()`` to grab the -current request's ``ServerSession`` and ``ChannelSink.bind`` it. - -Semantics: - -* **Last-claim-wins.** A second client calling ``claim_channel`` silently - replaces the previous binding; the prior leader stops receiving events. -* **Lossy on leader loss.** If the bound session goes away, the next - ``send_message`` raises and ``ChannelSink`` swallows it as a warning. - Events drop until another client claims. -* **stdio = trivially the one client.** Under stdio there is exactly one - session ever; calling ``claim_channel`` once after init binds it for - the rest of the server's life. Until then, channel events drop. -""" - -from ..base_step import BaseStep -from ...components import R - - -@R.register("claim_channel_step") -class ClaimChannelStep(BaseStep): - """Bind the current MCP session as the ```` recipient.""" - - async def execute(self): - if self.context is None: - raise RuntimeError("claim_channel_step requires 'context'") - - try: - from fastmcp.server.dependencies import get_context - - ctx = get_context() - session = ctx.session - if session is None: - raise RuntimeError("FastMCP context has no ServerSession") - if self.app_context is None: - raise RuntimeError("claim_channel requires an application context") - sink = self.app_context.metadata.get("channel_sink") - if sink is None: - raise RuntimeError("channel_sink not configured on application context metadata") - session_id = ctx.session_id or "" - sink.bind(session) - except Exception as e: - self.context.response.answer = {"claimed": False, "reason": f"{type(e).__name__}: {e}"} - self.context.response.metadata["claimed"] = False - return self.context.response - - self.logger.info(f"[claim_channel] channel bound to session={session_id}") - self.context.response.answer = { - "claimed": True, - "session_id": session_id, - "note": ( - "this session now receives notifications. " - "last-claim-wins: another call to claim_channel takes over." - ), - } - self.context.response.metadata["claimed"] = True - return self.context.response diff --git a/reme/steps/evolve/auto_memory_cc.py b/reme/steps/evolve/auto_memory_cc.py index b0ecd871..43b23219 100644 --- a/reme/steps/evolve/auto_memory_cc.py +++ b/reme/steps/evolve/auto_memory_cc.py @@ -15,7 +15,7 @@ abstraction and avoids the ``Msg`` round-trip entirely: defer to :class:`AutoMemoryStep` for the daily-note write/merge. Both the read (Claude Code side) and the copy (ReMe side) use the same -file-backed SessionStore, just rooted at different directories. +file-backed SessionStore with the SDK's project/session key layout. """ from __future__ import annotations @@ -52,6 +52,7 @@ class AutoMemoryCCStep(AutoMemoryStep): # Sub-directory under the session dir holding ReMe's copy of CC transcripts. _CC_STORE_SUBDIR = "claude_code" + _REME_PROJECT_KEY = "claude_code" async def execute(self): assert self.context is not None @@ -81,8 +82,8 @@ class AutoMemoryCCStep(AutoMemoryStep): transcript_dir = self._resolve_transcript_dir(session_id) if transcript_dir is None: return [] - store = CcFileSessionStore(transcript_dir) - return await store.load({"session_id": session_id}) or [] + store = CcFileSessionStore(self._projects_dir()) + return await store.load({"project_key": transcript_dir.name, "session_id": session_id}) or [] async def _save_cc_session(self, session_id: str, cc_entries: list[dict]) -> list[dict]: """Copy raw CC entries into ReMe's CC SessionStore; return the increment. @@ -96,7 +97,7 @@ class AutoMemoryCCStep(AutoMemoryStep): if not session_id: return [] store = self._reme_cc_store() - key = {"session_id": session_id} + key = {"project_key": self._REME_PROJECT_KEY, "session_id": session_id} cc_entries = [e for e in cc_entries if isinstance(e, dict) and e.get("uuid")] existing = await store.load(key) or [] seen = {e.get("uuid") for e in existing if isinstance(e, dict) and e.get("uuid")} @@ -105,7 +106,7 @@ class AutoMemoryCCStep(AutoMemoryStep): return increment def _reme_cc_store(self) -> CcFileSessionStore: - root = self.file_store.workspace_path / self._session_dir() / self._CC_STORE_SUBDIR + root = self.file_store.workspace_path / self._session_dir() return CcFileSessionStore(root) @staticmethod diff --git a/reme/utils/logger_utils.py b/reme/utils/logger_utils.py index 25d8a479..9d9afa07 100644 --- a/reme/utils/logger_utils.py +++ b/reme/utils/logger_utils.py @@ -68,7 +68,7 @@ def _init_loguru(log_dir: str, level: str, log_to_console: bool, log_to_file: bo try: os.makedirs(log_dir, exist_ok=True) current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - log_filepath = os.path.join(log_dir, f"{current_ts}.log") + log_filepath = os.path.join(log_dir, f"{current_ts}_{os.getpid()}.log") logger.add( log_filepath, @@ -116,6 +116,7 @@ def _init_stdlib(log_dir: str, level: str, log_to_console: bool, log_to_file: bo if log_to_file: try: os.makedirs(log_dir, exist_ok=True) + # Keep stdlib file naming aligned with QwenPaw logging. current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") log_filepath = os.path.join(log_dir, f"{current_ts}.log") diff --git a/tests/unit/test_auto_memory_cc.py b/tests/unit/test_auto_memory_cc.py new file mode 100644 index 00000000..46deeece --- /dev/null +++ b/tests/unit/test_auto_memory_cc.py @@ -0,0 +1,29 @@ +"""Unit tests for Claude Code auto-memory session persistence.""" + +# pylint: disable=protected-access + +from types import SimpleNamespace + +import pytest + +from reme.steps.evolve.auto_memory_cc import AutoMemoryCCStep + + +@pytest.mark.asyncio +async def test_reme_cc_store_preserves_existing_session_layout(tmp_path): + """Existing transcript UUIDs remain visible at session/claude_code/.jsonl.""" + step = AutoMemoryCCStep() + step.file_store = SimpleNamespace(workspace_path=tmp_path) + session_id = "session-1" + session_path = tmp_path / "session" / "claude_code" / f"{session_id}.jsonl" + session_path.parent.mkdir(parents=True) + session_path.write_text('{"uuid":"existing"}\n', encoding="utf-8") + + increment = await step._save_cc_session( + session_id, + [{"uuid": "existing"}, {"uuid": "new"}], + ) + + assert increment == [{"uuid": "new"}] + assert step._session_link(session_id) == f"[[session/claude_code/{session_id}.jsonl]]" + assert not (tmp_path / "session" / "claude_code" / "claude_code").exists() diff --git a/tests/unit/test_base_agent_wrapper.py b/tests/unit/test_base_agent_wrapper.py new file mode 100644 index 00000000..b76798bb --- /dev/null +++ b/tests/unit/test_base_agent_wrapper.py @@ -0,0 +1,57 @@ +"""Tests for shared agent wrapper behavior.""" + +from unittest.mock import MagicMock + +import pytest + +from reme.components.agent_wrapper import AsAgentWrapper, BaseAgentWrapper, CcAgentWrapper, CodexAgentWrapper +from reme.components.agent_wrapper import base_agent_wrapper +from reme.components import base_component + + +class _VersionedAgentWrapper(BaseAgentWrapper): + SDK_PACKAGE = "example-agent-sdk" + + async def reply(self, inputs, **kwargs) -> dict: + return {"inputs": inputs, "kwargs": kwargs} + + +def test_init_logs_sdk_version(monkeypatch): + """An SDK-backed wrapper logs its installed distribution version.""" + logger = MagicMock() + logger.bind.return_value = logger + monkeypatch.setattr(base_component, "get_logger", lambda: logger) + monkeypatch.setattr(base_agent_wrapper.metadata, "version", lambda package: "1.2.3") + + _VersionedAgentWrapper(name="versioned") + + logger.info.assert_called_once_with("Agent SDK package=example-agent-sdk version=1.2.3") + + +def test_init_logs_unknown_when_sdk_distribution_metadata_is_missing(monkeypatch): + """Missing distribution metadata does not prevent wrapper initialization.""" + logger = MagicMock() + logger.bind.return_value = logger + monkeypatch.setattr(base_component, "get_logger", lambda: logger) + + def missing_version(package): + raise base_agent_wrapper.metadata.PackageNotFoundError(package) + + monkeypatch.setattr(base_agent_wrapper.metadata, "version", missing_version) + + _VersionedAgentWrapper() + + logger.info.assert_called_once_with("Agent SDK package=example-agent-sdk version=unknown") + + +@pytest.mark.parametrize( + ("wrapper_class", "sdk_package"), + [ + (AsAgentWrapper, "agentscope"), + (CcAgentWrapper, "claude-agent-sdk"), + (CodexAgentWrapper, "openai-codex"), + ], +) +def test_agent_wrappers_declare_sdk_package(wrapper_class, sdk_package): + """Each concrete backend identifies the distribution that provides its SDK.""" + assert wrapper_class.SDK_PACKAGE == sdk_package diff --git a/tests/unit/test_base_component.py b/tests/unit/test_base_component.py index abaa7b50..6b9d994b 100644 --- a/tests/unit/test_base_component.py +++ b/tests/unit/test_base_component.py @@ -46,17 +46,6 @@ class FailCloseComponent(StubComponent): # -- Dependency --------------------------------------------------------------- -def test_dependency_repr_optional(): - dep = Dependency(ComponentEnum.FILE_CHUNKER, "my_parser", optional=True) - assert "?" in repr(dep) - assert "file_chunker" in repr(dep) - - -def test_dependency_repr_required(): - dep = Dependency(ComponentEnum.FILE_CHUNKER, "my_parser", optional=False) - assert "?" not in repr(dep) - - def test_dependency_getattr_raises(): dep = Dependency(ComponentEnum.FILE_CHUNKER, "my_parser") with pytest.raises(RuntimeError, match="accessed before start"): @@ -343,8 +332,6 @@ def test_component_metadata_path(): if __name__ == "__main__": print("\n=== BaseComponent Tests ===") - test_dependency_repr_optional() - test_dependency_repr_required() test_dependency_getattr_raises() test_bind_returns_none_for_empty_name() test_bind_returns_dependency_placeholder() diff --git a/tests/unit/test_bm25_index_perf.py b/tests/unit/test_bm25_index_perf.py deleted file mode 100644 index 12378392..00000000 --- a/tests/unit/test_bm25_index_perf.py +++ /dev/null @@ -1,337 +0,0 @@ -"""BM25Index performance tests for add_docs and retrieve.""" - -import asyncio -import os -import random -import tempfile -import time - -from reme.components.keyword_index import BM25Index -from reme.components.tokenizer import RegexTokenizer - -# A small vocab of realistic-looking words for generating random text -_VOCAB = [ - "algorithm", - "data", - "machine", - "learning", - "model", - "network", - "neural", - "training", - "optimization", - "gradient", - "loss", - "function", - "parameter", - "weight", - "bias", - "layer", - "activation", - "relu", - "sigmoid", - "softmax", - "backpropagation", - "forward", - "pass", - "batch", - "epoch", - "iteration", - "convergence", - "divergence", - "regularization", - "dropout", - "attention", - "transformer", - "encoder", - "decoder", - "embedding", - "token", - "vector", - "matrix", - "tensor", - "computation", - "graph", - "node", - "edge", - "vertex", - "path", - "search", - "retrieval", - "index", - "query", - "document", - "corpus", - "term", - "frequency", - "inverse", - "score", - "rank", - "relevance", - "precision", - "recall", - "f1", - "metric", - "evaluation", - "benchmark", - "dataset", - "sample", - "feature", - "label", - "class", - "predict", - "classification", - "regression", - "clustering", - "dimension", - "reduction", - "pca", - "tsne", - "visualization", - "matplotlib", - "plot", - "chart", - "histogram", - "scatter", - "line", - "bar", - "database", - "sql", - "query", - "table", - "row", - "column", - "index", - "primary", - "foreign", - "key", - "constraint", - "schema", - "migration", - "version", - "control", - "git", - "commit", - "branch", - "merge", - "conflict", - "resolution", - "review", - "approve", - "reject", - "pull", - "request", - "issue", - "bug", - "fix", - "feature", - "enhancement", - "refactor", - "test", - "deploy", - "production", - "staging", - "development", - "environment", - "configuration", - "setting", - "variable", - "constant", - "global", - "local", - "scope", - "closure", - "callback", - "promise", - "async", - "await", - "synchronous", - "asynchronous", - "concurrent", - "parallel", - "thread", - "process", - "memory", - "cache", - "buffer", - "queue", - "stack", - "heap", - "pool", -] - - -class temp_chdir: - """Context manager to temporarily chdir into a path and restore on exit.""" - - def __init__(self, path): - self.path = path - self.old = None - - def __enter__(self): - self.old = os.getcwd() - os.chdir(self.path) - return self - - def __exit__(self, *exc): - os.chdir(self.old) - - -def _gen_random_text(n_tokens: int) -> str: - """Generate random text with approximately n_tokens words.""" - words = random.choices(_VOCAB, k=n_tokens) - return " ".join(words) - - -def _gen_random_query(n_words: int) -> str: - """Generate a random query with n_words words.""" - words = random.choices(_VOCAB, k=n_words) - return " ".join(words) - - -async def _make_index() -> BM25Index: - """Create and start a BM25Index using cwd as working dir, with non-filtering tokenizer.""" - index = BM25Index() - tokenizer = RegexTokenizer(filter_stopwords=False) - index.tokenizer = tokenizer - index._owned.append(tokenizer) # pylint: disable=protected-access - await index.start() - return index - - -async def _setup_index_for_retrieve(n_docs: int = 100, doc_tokens: int = 1000) -> BM25Index: - """Build an index with n_docs medium-sized docs in cwd.""" - index = await _make_index() - docs = {f"doc_{i}": _gen_random_text(doc_tokens) for i in range(n_docs)} - await index.add_docs(docs) - return index - - -def test_add_docs_small(): - """Add 100 small docs (~100 tokens each).""" - - async def run(): - docs = {f"doc_{i}": _gen_random_text(100) for i in range(100)} - with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - index = await _make_index() - t0 = time.perf_counter() - await index.add_docs(docs) - elapsed = time.perf_counter() - t0 - print(f" add_docs (100 docs x ~100 tokens): {elapsed:.4f}s") - await index.close() - - asyncio.run(run()) - - -def test_add_docs_medium(): - """Add 100 medium docs (~1000 tokens each).""" - - async def run(): - docs = {f"doc_{i}": _gen_random_text(1000) for i in range(100)} - with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - index = await _make_index() - t0 = time.perf_counter() - await index.add_docs(docs) - elapsed = time.perf_counter() - t0 - print(f" add_docs (100 docs x ~1000 tokens): {elapsed:.4f}s") - await index.close() - - asyncio.run(run()) - - -def test_add_docs_large(): - """Add 100 large docs (~10000 tokens each).""" - - async def run(): - docs = {f"doc_{i}": _gen_random_text(10000) for i in range(100)} - with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - index = await _make_index() - t0 = time.perf_counter() - await index.add_docs(docs) - elapsed = time.perf_counter() - t0 - print(f" add_docs (100 docs x ~10000 tokens): {elapsed:.4f}s") - await index.close() - - asyncio.run(run()) - - -def test_retrieve_short_query(): - """Retrieve with 1-word query.""" - - async def run(): - with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - index = await _setup_index_for_retrieve() - query = _gen_random_query(1) - t0 = time.perf_counter() - await index.retrieve(query, limit=10) - elapsed = time.perf_counter() - t0 - print(f" retrieve (1-word query, 100 docs): {elapsed:.6f}s") - await index.close() - - asyncio.run(run()) - - -def test_retrieve_medium_query(): - """Retrieve with 5-word query.""" - - async def run(): - with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - index = await _setup_index_for_retrieve() - query = _gen_random_query(5) - t0 = time.perf_counter() - await index.retrieve(query, limit=10) - elapsed = time.perf_counter() - t0 - print(f" retrieve (5-word query, 100 docs): {elapsed:.6f}s") - await index.close() - - asyncio.run(run()) - - -def test_retrieve_long_query(): - """Retrieve with 20-word query.""" - - async def run(): - with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - index = await _setup_index_for_retrieve() - query = _gen_random_query(20) - t0 = time.perf_counter() - await index.retrieve(query, limit=10) - elapsed = time.perf_counter() - t0 - print(f" retrieve (20-word query, 100 docs): {elapsed:.6f}s") - await index.close() - - asyncio.run(run()) - - -def test_retrieve_very_long_query(): - """Retrieve with 100-word query.""" - - async def run(): - with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - index = await _setup_index_for_retrieve() - query = _gen_random_query(100) - t0 = time.perf_counter() - await index.retrieve(query, limit=10) - elapsed = time.perf_counter() - t0 - print(f" retrieve (100-word query, 100 docs): {elapsed:.6f}s") - await index.close() - - asyncio.run(run()) - - -if __name__ == "__main__": - random.seed(42) - print("=== BM25Index Performance Tests ===\n") - - print("[add_docs]") - test_add_docs_small() - test_add_docs_medium() - test_add_docs_large() - - print("\n[retrieve]") - test_retrieve_short_query() - test_retrieve_medium_query() - test_retrieve_long_query() - test_retrieve_very_long_query() - - print("\nDone.") diff --git a/tests/unit/test_cc_agent_wrapper.py b/tests/unit/test_cc_agent_wrapper.py index 671de056..9a84287e 100644 --- a/tests/unit/test_cc_agent_wrapper.py +++ b/tests/unit/test_cc_agent_wrapper.py @@ -1,14 +1,17 @@ """Tests for the Claude Code agent wrapper.""" from dataclasses import replace +from itertools import count from pathlib import Path +from types import SimpleNamespace import pytest from reme.components.agent_wrapper.as_agent_wrapper import AsAgentWrapper from reme.components.agent_wrapper.cc_agent_wrapper import CcAgentWrapper +from reme.components.agent_wrapper.cc_session_store import CcFileSessionStore from reme.components.application_context import ApplicationContext -from reme.config import resolve_app_config +from reme.enumeration import ChunkEnum, ComponentEnum # pylint: disable=protected-access @@ -24,7 +27,9 @@ def _skill_roots(tmp_path: Path) -> tuple[Path, Path]: ) -def test_ensure_claude_skill_dir_adds_selected_skills_without_replacing_existing(tmp_path): +def test_ensure_claude_skill_dir_adds_selected_skills_without_replacing_existing( + tmp_path, +): """Selected workspace skills are added while unrelated Claude skills remain.""" project_skills = tmp_path / "skills" (project_skills / "one").mkdir(parents=True) @@ -60,8 +65,8 @@ def test_ensure_claude_skill_dir_all_adds_each_project_skill(tmp_path): assert {path.name for path in root.iterdir()} == {"one", "two"} -def test_ensure_claude_skill_dir_migrates_old_directory_link(tmp_path): - """A legacy link to the whole project skills directory is migrated safely.""" +def test_ensure_claude_skill_dir_preserves_existing_directory_link(tmp_path): + """An existing skills link is user-owned and remains untouched.""" project_skills = tmp_path / "skills" (project_skills / "one").mkdir(parents=True) config_dir = tmp_path / "mem_session" / "claude_config" @@ -72,10 +77,19 @@ def test_ensure_claude_skill_dir_migrates_old_directory_link(tmp_path): _wrapper(tmp_path)._ensure_claude_skill_dir(config_dir, ["one"]) assert legacy_root.is_dir() - assert not legacy_root.is_symlink() + assert legacy_root.is_symlink() assert (legacy_root / "one").resolve() == (project_skills / "one").resolve() +@pytest.mark.asyncio +async def test_file_session_store_conforms_to_latest_sdk(tmp_path): + """The file store follows the SDK project/session/subkey contract.""" + from claude_agent_sdk.testing import run_session_store_conformance + + sequence = count() + await run_session_store_conformance(lambda: CcFileSessionStore(tmp_path / str(next(sequence)))) + + def test_ensure_claude_skill_dir_rejects_paths_as_skill_names(tmp_path): """Skill selectors cannot escape the project skills directory.""" (tmp_path / "skills").mkdir() @@ -84,23 +98,29 @@ def test_ensure_claude_skill_dir_rejects_paths_as_skill_names(tmp_path): _wrapper(tmp_path)._ensure_claude_skill_dir(tmp_path / "config", ["../outside"]) -def test_system_prompt_mode_replace_preserves_current_behavior(tmp_path): - """Replace mode passes a string system prompt directly to the SDK.""" +def test_configured_skills_use_latest_sdk_allowlist(tmp_path): + """Selected ReMe skills are passed through using the SDK's list semantics.""" + project_skills = tmp_path / "skills" + (project_skills / "one").mkdir(parents=True) + (project_skills / "two").mkdir() + + opts = _wrapper(tmp_path)._build_options("hello", skills=["one"]) + + assert opts.skills == ["one"] + for root in _skill_roots(tmp_path): + assert (root / "one").resolve() == (project_skills / "one").resolve() + assert not (root / "two").exists() + + +def test_sdk_native_system_prompt_preset_is_preserved(tmp_path): + """System prompt dictionaries pass directly to the latest SDK.""" opts = _wrapper(tmp_path)._build_options( "hello", - system_prompt="custom prompt", - system_prompt_mode="replace", - ) - - assert opts.system_prompt == "custom prompt" - - -def test_system_prompt_mode_append_uses_claude_code_preset(tmp_path): - """Append mode retains Claude Code's preset and appends the custom prompt.""" - opts = _wrapper(tmp_path)._build_options( - "hello", - system_prompt="custom prompt", - system_prompt_mode="append", + system_prompt={ + "type": "preset", + "preset": "claude_code", + "append": "custom prompt", + }, ) assert opts.system_prompt == { @@ -110,17 +130,42 @@ def test_system_prompt_mode_append_uses_claude_code_preset(tmp_path): } -def test_system_prompt_mode_rejects_unknown_value(tmp_path): - """Invalid prompt modes fail with a clear configuration error.""" - with pytest.raises(ValueError, match="Unknown system_prompt_mode"): - _wrapper(tmp_path)._build_options("hello", system_prompt_mode="merge") +def test_api_credentials_use_only_wrapper_config(tmp_path, monkeypatch): + """Claude Code credentials do not fall back to ambient or shared LLM configuration.""" + wrapper = _wrapper(tmp_path) + wrapper.app_context.app_config.environment = { + "ANTHROPIC_AUTH_TOKEN": "application-key", + "ANTHROPIC_BASE_URL": "https://application.example.test", + "TOOL_ENV": "preserved", + } + wrapper.app_context.app_config.components[ComponentEnum.AS_LLM] = { + "default": SimpleNamespace( + credential={ + "api_key": "default-key", + "base_url": "https://default.example.test", + }, + ), + } + for name in ("ANTHROPIC_AUTH_TOKEN", "CLAUDE_CODE_API_KEY", "LLM_API_KEY"): + monkeypatch.setenv(name, "ambient-key") + for name in ("ANTHROPIC_BASE_URL", "CLAUDE_CODE_BASE_URL", "LLM_BASE_URL"): + monkeypatch.setenv(name, "https://ambient.example.test") + configured = wrapper._build_options( # pylint: disable=protected-access + "hello", + api_key="configured-key", + base_url="https://configured.example.test", + credential={"api_key": "nested-key", "base_url": "https://nested.example.test"}, + ) + assert configured.env["ANTHROPIC_AUTH_TOKEN"] == "configured-key" + assert configured.env["ANTHROPIC_BASE_URL"] == "https://configured.example.test" + assert configured.env["TOOL_ENV"] == "preserved" - -def test_default_claude_code_system_prompt_mode_is_replace(): - """The built-in configuration preserves the existing replacement behavior.""" - config = resolve_app_config(log_config=False) - - assert config["components"]["agent_wrapper"]["claude_code"]["system_prompt_mode"] == "replace" + empty = wrapper._build_options( # pylint: disable=protected-access + "hello", + credential={"api_key": "nested-key", "base_url": "https://nested.example.test"}, + ) + assert empty.env["ANTHROPIC_AUTH_TOKEN"] == "" + assert empty.env["ANTHROPIC_BASE_URL"] == "" def test_build_options_accepts_empty_output_schema(tmp_path): @@ -130,6 +175,64 @@ def test_build_options_accepts_empty_output_schema(tmp_path): assert opts.output_format == {"type": "json_schema", "schema": {}} +def test_build_options_uses_native_sessions_and_allows_file_checkpointing(tmp_path): + """Local Claude transcripts remain the default and do not conflict with checkpoints.""" + opts = _wrapper(tmp_path)._build_options("hello", enable_file_checkpointing=True) + + assert opts.enable_file_checkpointing is True + assert opts.session_store is None + assert opts.env["CLAUDE_CONFIG_DIR"] == str(tmp_path / "mem_session" / "claude_config") + + +def test_build_options_preserves_explicit_session_store(tmp_path): + """Callers can still opt into the SDK's external transcript mirror.""" + store = CcFileSessionStore(tmp_path / "mirror") + + opts = _wrapper(tmp_path)._build_options("hello", session_store=store) + + assert opts.session_store is store + + +def test_job_tools_reject_non_mapping_mcp_config_instead_of_discarding_it(tmp_path, monkeypatch): + """Adding ReMe tools never silently replaces an SDK MCP config path.""" + wrapper = _wrapper(tmp_path) + job = SimpleNamespace(name="remember", description="Remember", parameters={}) + monkeypatch.setattr(wrapper, "_resolve_job_tools", lambda _names: [job]) + + with pytest.raises(ValueError, match="mcp_servers to be a mapping"): + wrapper._build_options("hello", job_tools=["remember"], mcp_servers=tmp_path / "mcp.json") + + +def test_job_tools_merge_with_mapping_mcp_config(tmp_path, monkeypatch): + """Existing MCP configuration remains reusable beside ReMe tools.""" + wrapper = _wrapper(tmp_path) + job = SimpleNamespace(name="remember", description="Remember", parameters={}) + monkeypatch.setattr(wrapper, "_resolve_job_tools", lambda _names: [job]) + external = {"type": "http", "url": "https://mcp.example.test"} + mcp_servers = {"external": external} + allowed_tools = ["Read"] + + first = wrapper._build_options( + "hello", + job_tools=["remember"], + mcp_servers=mcp_servers, + allowed_tools=allowed_tools, + ) + second = wrapper._build_options( + "hello", + job_tools=["remember"], + mcp_servers=mcp_servers, + allowed_tools=allowed_tools, + ) + + assert mcp_servers == {"external": external} + assert allowed_tools == ["Read"] + for opts in (first, second): + assert opts.mcp_servers["external"] is external + assert opts.mcp_servers[wrapper.MCP_SERVER_NAME]["type"] == "sdk" + assert opts.allowed_tools == ["Read", "remember"] + + @pytest.mark.asyncio async def test_reply_preserves_falsy_structured_output(tmp_path, monkeypatch): """Falsy structured output is returned instead of being discarded.""" @@ -157,6 +260,206 @@ async def test_reply_preserves_falsy_structured_output(tmp_path, monkeypatch): assert result["structured_output"] == {} +def test_error_result_with_success_subtype_is_not_suppressed(): + """Latest SDK can report API failures with subtype=success and is_error=True.""" + from claude_agent_sdk import ResultMessage + + message = ResultMessage( + subtype="success", + duration_ms=1, + duration_api_ms=1, + is_error=True, + num_turns=1, + session_id="session-1", + errors=["upstream unavailable"], + api_error_status=529, + ) + + chunks = CcAgentWrapper._result_message_to_chunks(message) + + error = next(chunk for chunk in chunks if chunk.chunk_type.value == "error") + assert error.chunk == "upstream unavailable" + assert error.metadata == {"api_error_status": 529} + + +def test_latest_sdk_server_tool_blocks_are_converted(): + """Server-side tools use the same unified call/result lifecycle.""" + from claude_agent_sdk import AssistantMessage, ServerToolResultBlock + + call = CcAgentWrapper._raw_event_to_chunk( + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "server_tool_use", + "id": "tool-1", + "name": "web_search", + }, + }, + ) + message = AssistantMessage( + content=[ServerToolResultBlock(tool_use_id="tool-1", content={"type": "web_search_result"})], + model="claude", + ) + results = CcAgentWrapper._message_content_to_chunks(message, visible_tool_call_ids={"tool-1"}) + + assert call is not None and call.chunk_type.value == "tool_call" + assert len(results) == 1 + assert results[0].chunk_type.value == "tool_result" + assert results[0].chunk == { + "tool_use_id": "tool-1", + "content": {"type": "web_search_result"}, + } + + +@pytest.mark.asyncio +async def test_reply_stream_emits_one_reply_end_for_normal_sdk_lifecycle(tmp_path, monkeypatch): + """Message delta reports usage while message stop is the sole lifecycle end.""" + from claude_agent_sdk import ResultMessage, StreamEvent + + async def query(**_kwargs): + for event in ( + { + "type": "message_start", + "message": {"id": "message-1", "role": "assistant"}, + }, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 2}, + }, + {"type": "message_stop"}, + ): + yield StreamEvent(uuid="event-1", session_id="session-1", event=event) + yield ResultMessage( + subtype="success", + duration_ms=1, + duration_api_ms=1, + is_error=False, + num_turns=1, + session_id="session-1", + usage={"input_tokens": 1, "output_tokens": 2}, + ) + + monkeypatch.setattr("claude_agent_sdk.query", query) + chunks = [chunk async for chunk in _wrapper(tmp_path).reply_stream("hello")] + + assert sum(chunk.chunk_type == ChunkEnum.REPLY_END for chunk in chunks) == 1 + delta_usage = next(chunk for chunk in chunks if chunk.metadata.get("stop_reason") == "end_turn") + assert delta_usage.chunk_type == ChunkEnum.USAGE + assert delta_usage.output_tokens == 2 + + +@pytest.mark.asyncio +async def test_reply_stream_only_reports_rejected_rate_limit(tmp_path, monkeypatch): + """Rate-limit warnings are informational; only rejected is an error.""" + from claude_agent_sdk import RateLimitEvent, RateLimitInfo, ResultMessage + + async def query(**_kwargs): + yield RateLimitEvent( + rate_limit_info=RateLimitInfo(status="allowed_warning"), + uuid="warning", + session_id="session-1", + ) + yield RateLimitEvent( + rate_limit_info=RateLimitInfo(status="rejected"), + uuid="rejected", + session_id="session-1", + ) + yield ResultMessage( + subtype="success", + duration_ms=1, + duration_api_ms=1, + is_error=False, + num_turns=1, + session_id="session-1", + ) + + monkeypatch.setattr("claude_agent_sdk.query", query) + chunks = [chunk async for chunk in _wrapper(tmp_path).reply_stream("hello")] + + errors = [chunk for chunk in chunks if chunk.chunk_type.value == "error"] + assert [chunk.chunk for chunk in errors] == ["Rate limit exceeded"] + + +@pytest.mark.asyncio +async def test_reply_stream_uses_error_result_as_terminal_error(tmp_path, monkeypatch): + """The SDK's non-zero process exit does not duplicate an emitted error result.""" + from claude_agent_sdk import ResultMessage + + async def query(**_kwargs): + yield ResultMessage( + subtype="error_during_execution", + duration_ms=1, + duration_api_ms=1, + is_error=True, + num_turns=1, + session_id="session-1", + errors=["tool failed"], + ) + raise RuntimeError("Claude Code returned an error result: tool failed") + + monkeypatch.setattr("claude_agent_sdk.query", query) + chunks = [chunk async for chunk in _wrapper(tmp_path).reply_stream("hello")] + + errors = [chunk for chunk in chunks if chunk.chunk_type.value == "error"] + assert [chunk.chunk for chunk in errors] == ["tool failed"] + + +@pytest.mark.asyncio +async def test_reply_stream_does_not_hide_unrelated_error_after_error_result(tmp_path, monkeypatch): + """Only the SDK's exact trailing process error is suppressed.""" + from claude_agent_sdk import ResultMessage + + async def query(**_kwargs): + yield ResultMessage( + subtype="error_during_execution", + duration_ms=1, + duration_api_ms=1, + is_error=True, + num_turns=1, + session_id="session-1", + errors=["tool failed"], + ) + raise RuntimeError("unrelated store failure") + + monkeypatch.setattr("claude_agent_sdk.query", query) + + with pytest.raises(RuntimeError, match="unrelated store failure"): + _ = [chunk async for chunk in _wrapper(tmp_path).reply_stream("hello")] + + +@pytest.mark.asyncio +async def test_reply_stream_reports_session_mirror_errors(tmp_path, monkeypatch): + """The latest SDK's non-fatal mirror failures remain visible to callers.""" + from claude_agent_sdk import MirrorErrorMessage, ResultMessage + + async def query(**_kwargs): + yield MirrorErrorMessage( + subtype="mirror_error", + data={}, + key={"project_key": "project", "session_id": "session-1"}, + error="disk full", + ) + yield ResultMessage( + subtype="success", + duration_ms=1, + duration_api_ms=1, + is_error=False, + num_turns=1, + session_id="session-1", + ) + + monkeypatch.setattr("claude_agent_sdk.query", query) + chunks = [chunk async for chunk in _wrapper(tmp_path).reply_stream("hello")] + + diagnostics = [chunk for chunk in chunks if chunk.metadata.get("event") == "session_mirror_error"] + assert len(diagnostics) == 1 + assert diagnostics[0].chunk_type == ChunkEnum.DATA + assert diagnostics[0].chunk == "Session mirror failed: disk full" + assert not any(chunk.chunk_type == ChunkEnum.ERROR for chunk in chunks) + + @pytest.mark.asyncio @pytest.mark.parametrize( "wrapper_factory", diff --git a/tests/unit/test_channel_notify.py b/tests/unit/test_channel_notify.py deleted file mode 100644 index c78aa6a9..00000000 --- a/tests/unit/test_channel_notify.py +++ /dev/null @@ -1,107 +0,0 @@ -"""Tests for ``ChannelNotifyStep`` — workspace-watcher batch → channel event.""" - -import asyncio -from pathlib import Path - -from reme.components.application_context import ApplicationContext -from reme.components.service.mcp_service import ChannelSink -from reme.components.runtime_context import RuntimeContext -from reme.steps.channel.channel_notify import ChannelNotifyStep - - -class _StubSession: - """Capture ``send_message`` payloads instead of writing them to a transport.""" - - def __init__(self) -> None: - self.sent: list = [] - - async def send_message(self, message) -> None: - """Record the outbound ``SessionMessage`` for later assertions.""" - self.sent.append(message) - - -def _run(coro): - """Drive a coroutine on a fresh event loop (tests don't share one).""" - return asyncio.new_event_loop().run_until_complete(coro) - - -def _ctx(changes: list[dict]) -> RuntimeContext: - """Build a ``RuntimeContext`` pre-populated with the step's ``changes`` input.""" - ctx = RuntimeContext() - ctx["changes"] = changes - return ctx - - -def _app_ctx_with_sink(workspace: Path, stub: _StubSession | None) -> tuple[ApplicationContext, ChannelSink | None]: - """Build an ``ApplicationContext`` rooted at ``workspace``; attach a sink bound to ``stub`` if given.""" - app_ctx = ApplicationContext(workspace_dir=str(workspace), app_name="reme-test") - if stub is None: - return app_ctx, None - sink = ChannelSink() - sink.bind(stub) - app_ctx.metadata["channel_sink"] = sink - return app_ctx, sink - - -def test_emits_one_event_per_batch_with_relative_paths(tmp_path): - """A batch of changes → exactly one notification with relative paths and a count meta.""" - workspace = tmp_path - (workspace / "resource" / "2026-06-03").mkdir(parents=True) - f1 = workspace / "resource" / "2026-06-03" / "a.md" - f1.write_text("x") - - stub = _StubSession() - app_ctx, _ = _app_ctx_with_sink(workspace, stub) - - step = ChannelNotifyStep(app_context=app_ctx) - response = _run( - step( - context=_ctx( - [ - {"change": "added", "path": str(f1)}, - {"change": "modified", "path": str(f1)}, - ], - ), - ), - ) - - assert response.success is True - assert len(stub.sent) == 1 - params = stub.sent[0].message.root.params - assert params["meta"] == {"kind": "workspace_change", "count": "2"} - assert "added: resource/2026-06-03/a.md" in params["content"] - assert "modified: resource/2026-06-03/a.md" in params["content"] - - -def test_noop_when_no_changes(tmp_path): - """An empty changes list must not produce any notification.""" - stub = _StubSession() - app_ctx, _ = _app_ctx_with_sink(tmp_path, stub) - step = ChannelNotifyStep(app_context=app_ctx) - response = _run(step(context=_ctx([]))) - assert response.success is True - assert not stub.sent - - -def test_noop_when_sink_not_bound(tmp_path): - """Step must run cleanly when no ``ChannelSink`` is configured.""" - app_ctx, _ = _app_ctx_with_sink(tmp_path, None) - step = ChannelNotifyStep(app_context=app_ctx) - # Should run without raising even though channel_sink is absent from metadata - response = _run(step(context=_ctx([{"change": "added", "path": "/tmp/x.md"}]))) - assert response.success is True - - -def test_path_outside_workspace_passes_through_as_is(tmp_path): - """Paths not under the workspace are emitted verbatim instead of crashing.""" - stub = _StubSession() - app_ctx, _ = _app_ctx_with_sink(tmp_path, stub) - step = ChannelNotifyStep(app_context=app_ctx) - _run( - step( - context=_ctx([{"change": "added", "path": "/elsewhere/wild.md"}]), - ), - ) - # Stray absolute path → emitted verbatim, no crash, still one event - assert len(stub.sent) == 1 - assert "added: /elsewhere/wild.md" in stub.sent[0].message.root.params["content"] diff --git a/tests/unit/test_channel_sink.py b/tests/unit/test_channel_sink.py deleted file mode 100644 index 438f477d..00000000 --- a/tests/unit/test_channel_sink.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Tests for ``ChannelSink`` — outbound ``notifications/claude/channel`` plumbing. - -Strategy: stub a session-like object with an async ``send_message`` capture -list and exercise three paths: - -* not bound → emit is a no-op (no exception, no captured message) -* bound + valid meta → captured JSON-RPC notification carries our method - + content + the meta we passed -* bound + meta with non-identifier keys → the bad keys are dropped, the - rest passes through verbatim -""" - -import asyncio - -from reme.components.service.mcp_service import ChannelSink - - -class _StubSession: - """Capture ``send_message`` payloads instead of writing them to a transport.""" - - def __init__(self) -> None: - self.sent: list = [] - - async def send_message(self, message) -> None: - """Record the outbound ``SessionMessage`` for later assertions.""" - self.sent.append(message) - - -def _run(coro): - """Drive a coroutine on a fresh event loop (tests don't share one).""" - return asyncio.new_event_loop().run_until_complete(coro) - - -def test_emit_without_bind_is_noop(): - """Emitting before any session is bound must silently no-op.""" - sink = ChannelSink() - _run(sink.emit("hello", {"k": "v"})) # must not raise - - -def test_emit_after_bind_sends_channel_notification(): - """A bound session receives a JSON-RPC notification carrying content + meta verbatim.""" - sink = ChannelSink() - stub = _StubSession() - sink.bind(stub) - - _run(sink.emit("ingest done", {"path": "resource/2026-06-03/x.md", "kind": "ingest"})) - - assert len(stub.sent) == 1 - payload = stub.sent[0].message.root - assert payload.method == "notifications/claude/channel" - assert payload.params["content"] == "ingest done" - assert payload.params["meta"] == {"path": "resource/2026-06-03/x.md", "kind": "ingest"} - - -def test_emit_filters_non_identifier_meta_keys(): - """Meta keys that aren't pure ``[A-Za-z0-9_]`` identifiers are dropped before send.""" - sink = ChannelSink() - stub = _StubSession() - sink.bind(stub) - - _run( - sink.emit( - "x", - { - "good_key": "ok", - "bad-key": "dropped", # hyphen - "also.bad": "dropped", # dot - "Number9": "kept", - }, - ), - ) - - meta = stub.sent[0].message.root.params["meta"] - assert meta == {"good_key": "ok", "Number9": "kept"} - - -def test_unbind_returns_to_noop(): - """After ``unbind``, subsequent emits stop reaching the previously bound session.""" - sink = ChannelSink() - stub = _StubSession() - sink.bind(stub) - sink.unbind() - - _run(sink.emit("x", {})) - assert not stub.sent - - -def test_emit_swallows_send_failures(): - """A failing send_message must not bubble out (notification is best-effort).""" - - class _BoomSession: - """Session whose ``send_message`` always raises, to exercise the failure path.""" - - async def send_message(self, message): - """Raise to simulate a broken transport.""" - raise RuntimeError("transport broke") - - sink = ChannelSink() - sink.bind(_BoomSession()) - _run(sink.emit("x", {})) # must not raise diff --git a/tests/unit/test_claim_channel.py b/tests/unit/test_claim_channel.py deleted file mode 100644 index f8b04ff8..00000000 --- a/tests/unit/test_claim_channel.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Tests for ``ClaimChannelStep`` — current MCP session binding.""" - -import asyncio -import subprocess -import sys -from types import SimpleNamespace - -from fastmcp.server.context import _current_context - -from reme.components.application_context import ApplicationContext -from reme.components.service.mcp_service import ChannelSink -from reme.steps.channel.claim_channel import ClaimChannelStep - - -class _StubSession: - """Capture outbound channel messages after being bound.""" - - def __init__(self) -> None: - self.sent: list = [] - - async def send_message(self, message) -> None: - """Record a message sent by the channel sink.""" - self.sent.append(message) - - -def _run(coro): - """Drive a coroutine on a fresh event loop.""" - return asyncio.new_event_loop().run_until_complete(coro) - - -def test_claim_channel_binds_current_session(tmp_path): - """The active FastMCP session becomes the sink recipient.""" - app_ctx = ApplicationContext(workspace_dir=str(tmp_path), app_name="reme-test") - sink = ChannelSink() - app_ctx.metadata["channel_sink"] = sink - session = _StubSession() - ctx = SimpleNamespace(session=session, session_id="sid-1") - token = _current_context.set(ctx) - try: - step = ClaimChannelStep(app_context=app_ctx) - response = _run(step()) - finally: - _current_context.reset(token) - - assert response.success is True - assert response.answer["claimed"] is True - assert response.answer["session_id"] == "sid-1" - assert response.metadata["claimed"] is True - - _run(sink.emit("hello", {"kind": "test"})) - assert len(session.sent) == 1 - - -def test_claim_channel_reports_missing_fastmcp_context(tmp_path): - """Calling outside a FastMCP request reports a clean unclaimed result.""" - app_ctx = ApplicationContext(workspace_dir=str(tmp_path), app_name="reme-test") - app_ctx.metadata["channel_sink"] = ChannelSink() - step = ClaimChannelStep(app_context=app_ctx) - - response = _run(step()) - - assert response.success is True - assert response.answer["claimed"] is False - assert response.metadata["claimed"] is False - assert "No active context" in response.answer["reason"] - - -def test_claim_channel_missing_sink_is_still_controlled_under_optimized_python(tmp_path): - """Runtime validation must not rely on assert, which Python -O removes.""" - code = f""" -import asyncio -from types import SimpleNamespace -from fastmcp.server.context import _current_context -from reme.components.application_context import ApplicationContext -from reme.steps.channel.claim_channel import ClaimChannelStep - -class Session: - async def send_message(self, message): - pass - -async def main(): - app_ctx = ApplicationContext(workspace_dir={str(tmp_path)!r}, app_name="reme-test") - ctx = SimpleNamespace(session=Session(), session_id="sid-optimized") - token = _current_context.set(ctx) - try: - response = await ClaimChannelStep(app_context=app_ctx)() - print(response.answer) - finally: - _current_context.reset(token) - -asyncio.run(main()) -""" - result = subprocess.run( - [sys.executable, "-O", "-c", code], - check=True, - capture_output=True, - text=True, - ) - assert "'claimed': False" in result.stdout - assert "channel_sink not configured" in result.stdout diff --git a/tests/unit/test_codex_agent_wrapper.py b/tests/unit/test_codex_agent_wrapper.py index 43d65e13..7d153c94 100644 --- a/tests/unit/test_codex_agent_wrapper.py +++ b/tests/unit/test_codex_agent_wrapper.py @@ -1,6 +1,6 @@ """Unit tests for the Codex agent wrapper and its FastMCP bridge.""" -# pylint: disable=missing-class-docstring,missing-function-docstring,protected-access +# pylint: disable=missing-class-docstring,missing-function-docstring,protected-access,too-many-lines import asyncio from dataclasses import dataclass @@ -15,7 +15,7 @@ from openai_codex.generated.v2_all import TokenUsageBreakdown from pydantic import BaseModel from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper -from reme.components.agent_wrapper.codex_mcp_server import _load_job_names, _make_tool, build_server +from reme.components.agent_wrapper.codex_mcp_server import _prepare_config from reme.components.job import BackgroundJob from reme.config import resolve_app_config from reme.enumeration import ChunkEnum, ComponentEnum @@ -43,6 +43,7 @@ def _wrapper(tmp_path, **kwargs): config = SimpleNamespace( workspace_dir=str(tmp_path), mem_session_dir="mem_session", + environment={}, components={ComponentEnum.AS_LLM: {}}, model_dump=lambda **_kwargs: { "workspace_dir": str(tmp_path), @@ -61,11 +62,13 @@ def test_mcp_config_uses_stdio_bridge_and_selected_jobs(tmp_path): wrapper, _job = _wrapper(tmp_path, mcp_config="custom.yaml") config = wrapper._mcp_server_config( # pylint: disable=protected-access - {"job_tools": ["search"], "tool_context_id": "ctx-1"}, + {"job_tools": ["search", "search"], "tool_context_id": "ctx-1"}, ) assert config["command"] assert config["enabled_tools"] == ["search"] + assert config["args"].count("--job") == 1 + assert config["args"][config["args"].index("--job") + 1] == "search" assert "reme.components.agent_wrapper.codex_mcp_server" in config["args"] assert config["args"][config["args"].index("--config") + 1] == str(tmp_path / "custom.yaml") assert config["args"][config["args"].index("--tool-context-id") + 1] == "ctx-1" @@ -93,46 +96,36 @@ def test_mcp_config_rejects_background_jobs(tmp_path): wrapper._mcp_server_config({"job_tools": ["watch"]}) -def test_bridge_tool_injects_tool_context_id(): - async def run(): - job = _Job() - tool = _make_tool(job, "ctx-1") - result = await tool.run({"query": "alpha"}) - assert job.calls == [{"query": "alpha", "tool_context_id": "ctx-1"}] - assert "found:alpha" in str(result.content) - - asyncio.run(run()) - - -def test_bridge_rejects_caller_tool_context_id(): - async def run(): - job = _Job() - tool = _make_tool(job, "ctx-1") - with pytest.raises(Exception, match="managed by the Codex agent wrapper"): - await tool.run({"query": "alpha", "tool_context_id": "caller"}) - - asyncio.run(run()) - - -def test_build_server_registers_only_selected_jobs(): - app = SimpleNamespace( - context=SimpleNamespace(jobs={"one": _Job("one"), "two": _Job("two")}), - start=lambda: None, - close=lambda: None, +def test_prepare_config_reuses_selected_stdio_mcp_service(): + prepared = _prepare_config( + { + "service": {"backend": "http"}, + "jobs": { + "selected": {"backend": "base", "enable_serve": False}, + "helper": {"backend": "base"}, + "watch": {"backend": "background"}, + }, + }, + ["selected"], + "ctx-1", ) - async def run(): - server = build_server(app, ["two"]) - tools = await server.list_tools(run_middleware=False) - assert [tool.name for tool in tools] == ["two"] - - asyncio.run(run()) + assert prepared["service"] == { + "backend": "mcp", + "transport": "stdio", + "jobs": ["selected"], + "tool_error_on_failure": True, + "injected_job_kwargs": {"tool_context_id": "ctx-1"}, + } + assert set(prepared["jobs"]) == {"selected", "helper"} + assert prepared["jobs"]["selected"]["enable_serve"] is True -def test_load_job_names_validates_json_array(): - assert _load_job_names('["one", "two"]') == ["one", "two"] - with pytest.raises(ValueError, match="JSON array"): - _load_job_names('{"one": true}') +def test_prepare_config_rejects_missing_or_background_selected_jobs(): + with pytest.raises(KeyError, match="missing"): + _prepare_config({"jobs": {}}, ["missing"]) + with pytest.raises(KeyError, match="watch"): + _prepare_config({"jobs": {"watch": {"backend": "background"}}}, ["watch"]) @pytest.mark.asyncio @@ -167,8 +160,8 @@ async def test_stdio_bridge_starts_and_lists_selected_job(tmp_path): str(config_path), "--workspace", str(tmp_path / "workspace"), - "--jobs", - '["empty"]', + "--job", + "empty", ], cwd=str(Path(__file__).resolve().parents[2]), ) @@ -207,8 +200,8 @@ async def test_stdio_bridge_stdout_is_protocol_clean(tmp_path): str(config_path), "--workspace", str(tmp_path / "workspace"), - "--jobs", - '["empty"]', + "--job", + "empty", cwd=str(Path(__file__).resolve().parents[2]), stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, @@ -297,13 +290,38 @@ def test_event_to_chunks_maps_content_usage_and_completion(): assert completed[0].metadata["status"] == "completed" +def test_event_to_chunks_preserves_new_turn_scoped_notifications(): + from openai_codex.types import Notification + from openai_codex.generated.v2_all import TurnDiffUpdatedNotification + + event = Notification( + method="turn/diff/updated", + payload=TurnDiffUpdatedNotification( + threadId="thread-1", + turnId="turn-1", + diff="diff --git a/a b/a", + ), + ) + + chunk = CodexAgentWrapper._event_to_chunks(event, "thread-1")[0] # pylint: disable=protected-access + + assert chunk.chunk_type == ChunkEnum.DATA + assert chunk.chunk == { + "threadId": "thread-1", + "turnId": "turn-1", + "diff": "diff --git a/a b/a", + } + assert chunk.metadata == {"codex_method": "turn/diff/updated"} + + @dataclass class _TurnResult: final_response: str status: str = "completed" -def test_reply_returns_thread_id_and_structured_output(tmp_path, monkeypatch): +@pytest.mark.asyncio +async def test_reply_returns_thread_id_and_structured_output(tmp_path, monkeypatch): wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") class FakeThread: @@ -333,15 +351,52 @@ def test_reply_returns_thread_id_and_structured_output(tmp_path, monkeypatch): async def thread_start(self, **_kwargs): return FakeThread() - monkeypatch.setattr("openai_codex.AsyncCodex", FakeCodex) - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.load_env", lambda *_args: {}) + monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - result = asyncio.run(wrapper.reply("answer", output_schema={"type": "object"})) + result = await wrapper.reply("answer", output_schema={"type": "object"}) + await wrapper.close() assert result["session_id"] == "thread-1" assert result["structured_output"] == {"ok": True} +@pytest.mark.asyncio +async def test_reply_accepts_latest_sdk_run_input(tmp_path, monkeypatch): + from openai_codex import LocalImageInput, TextInput + + wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") + inputs = [TextInput("describe this image"), LocalImageInput("image.png")] + observed = {} + + class FakeThread: + id = "thread-1" + + async def run(self, run_input, **_kwargs): + observed["input"] = run_input + return _TurnResult(final_response="done") + + class FakeCodex: + def __init__(self, _config): + pass + + async def account(self): + return SimpleNamespace(account=SimpleNamespace()) + + async def close(self): + return None + + async def thread_start(self, **_kwargs): + return FakeThread() + + monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) + + result = await wrapper.reply(inputs) + await wrapper.close() + + assert observed["input"] is inputs + assert result["last_message"] == "done" + + def test_codex_skills_add_all_without_deleting_existing_content(tmp_path): wrapper, _job = _wrapper(tmp_path) for name in ("reme_memory", "qwenpaw_memory"): @@ -507,6 +562,53 @@ async def test_open_thread_defaults_to_full_access(tmp_path): assert observed["sandbox"] == Sandbox.full_access +@pytest.mark.asyncio +async def test_open_thread_forwards_latest_sdk_options(tmp_path): + from openai_codex.types import Personality, ThreadSource, ThreadStartSource + + wrapper, _job = _wrapper(tmp_path) + observed = {} + + class FakeCodex: + async def thread_start(self, **kwargs): + observed["start"] = kwargs + return SimpleNamespace(id="thread-1") + + async def thread_resume(self, _thread_id, **kwargs): + observed["resume"] = kwargs + return SimpleNamespace(id="thread-1") + + async def thread_fork(self, _thread_id, **kwargs): + observed["fork"] = kwargs + return SimpleNamespace(id="thread-2") + + codex = FakeCodex() + await wrapper._open_thread( # pylint: disable=protected-access + codex, + { + "personality": "friendly", + "service_name": "reme", + "session_start_source": "startup", + "thread_source": "user", + }, + ) + await wrapper._open_thread( # pylint: disable=protected-access + codex, + {"session_id": "thread-1", "personality": "pragmatic"}, + ) + await wrapper._open_thread( # pylint: disable=protected-access + codex, + {"session_id": "thread-1", "fork_session": True, "thread_source": "subagent"}, + ) + + assert observed["start"]["personality"] == Personality.friendly + assert observed["start"]["service_name"] == "reme" + assert observed["start"]["session_start_source"] == ThreadStartSource.startup + assert observed["start"]["thread_source"] == ThreadSource.user + assert observed["resume"]["personality"] == Personality.pragmatic + assert observed["fork"]["thread_source"] == ThreadSource.subagent + + @pytest.mark.asyncio async def test_resume_reuses_tool_context_and_rejects_context_change(tmp_path): wrapper, _job = _wrapper(tmp_path) @@ -585,8 +687,7 @@ async def test_reply_normalizes_schema_and_reuses_persistent_client(tmp_path, mo async def thread_start(self, **_kwargs): return FakeThread() - monkeypatch.setattr("openai_codex.AsyncCodex", FakeCodex) - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.load_env", lambda *_args: {}) + monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) await wrapper.start() result = await wrapper.reply("first", output_schema=_StructuredModel) @@ -609,7 +710,7 @@ async def test_reply_stream_rejects_output_schema(tmp_path, schema): @pytest.mark.asyncio async def test_reply_stream_interrupts_turn_when_consumer_closes_early(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path) + wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") stream_closed = False interrupt_count = 0 @@ -637,75 +738,182 @@ async def test_reply_stream_interrupts_turn_when_consumer_closes_early(tmp_path, async def turn(self, _inputs, **_kwargs): return FakeTurn() - async def get_codex(_kwargs): - return SimpleNamespace() + class FakeCodex: + def __init__(self, _config): + pass - async def open_thread(_codex, _kwargs): - return FakeThread() + async def account(self): + return SimpleNamespace(account=SimpleNamespace()) - monkeypatch.setattr(wrapper, "_get_codex", get_codex) - monkeypatch.setattr(wrapper, "_open_thread", open_thread) + async def close(self): + return None + + async def thread_start(self, **_kwargs): + return FakeThread() + + monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) stream = wrapper.reply_stream("answer") first = await anext(stream) assert first.chunk_type == ChunkEnum.REPLY_START await stream.aclose() + await wrapper.close() assert interrupt_count == 1 assert stream_closed @pytest.mark.asyncio -async def test_persistent_client_rejects_launch_config_changes(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path) +async def test_close_waits_for_active_turn(tmp_path, monkeypatch): + wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") + turn_started = asyncio.Event() + release_turn = asyncio.Event() + client_closed = asyncio.Event() + + class FakeThread: + id = "thread-1" + + async def run(self, _inputs, **_kwargs): + turn_started.set() + await release_turn.wait() + return _TurnResult(final_response="done") class FakeCodex: def __init__(self, _config): pass + async def account(self): + return SimpleNamespace(account=SimpleNamespace()) + + async def close(self): + client_closed.set() + + async def thread_start(self, **_kwargs): + return FakeThread() + + monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) + + await wrapper.start() + reply_task = asyncio.create_task(wrapper.reply("answer")) + await turn_started.wait() + close_task = asyncio.create_task(wrapper.close()) + await asyncio.sleep(0) + + assert not client_closed.is_set() + assert not close_task.done() + + release_turn.set() + result = await reply_task + await close_task + + assert result["last_message"] == "done" + assert client_closed.is_set() + + +@pytest.mark.asyncio +async def test_component_start_keeps_optional_client_lazy(tmp_path, monkeypatch): + wrapper, _job = _wrapper(tmp_path, auth_mode="api_key", api_key="") + + def fail_if_constructed(_config): + raise AssertionError("Codex client should be lazy") + + monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", fail_if_constructed) + + await wrapper.start() + await wrapper.close() + + assert wrapper._codex is None + + +@pytest.mark.asyncio +async def test_client_config_is_fixed_for_component_lifetime(tmp_path, monkeypatch): + wrapper, _job = _wrapper(tmp_path, auth_mode="api_key", api_key="one") + clients = [] + + class FakeCodex: + def __init__(self, _config): + clients.append(self) + async def login_api_key(self, _api_key): return None async def close(self): return None - monkeypatch.setattr("openai_codex.AsyncCodex", FakeCodex) - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.load_env", lambda *_args: {}) + monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) await wrapper.start() - first = await wrapper._get_codex({"api_key": "one"}) # pylint: disable=protected-access - assert await wrapper._get_codex({"api_key": "one"}) is first # pylint: disable=protected-access - with pytest.raises(RuntimeError, match="configuration changed"): - await wrapper._get_codex({"api_key": "two"}) # pylint: disable=protected-access + assert await wrapper._get_codex() is clients[0] # pylint: disable=protected-access + with pytest.raises(TypeError, match="configured on the wrapper: api_key"): + await wrapper.reply("answer", api_key="two") await wrapper.close() + assert len(clients) == 1 + + +@pytest.mark.parametrize( + ("name", "value"), + [ + ("auth_mode", "oauth"), + ("base_url", "https://example.test/v1"), + ("codex_bin", "/tmp/codex"), + ("codex_home", "/tmp/codex-home"), + ("config_overrides", ['model="test"']), + ("cwd", "/tmp"), + ("experimental_api", False), + ("launch_args_override", ["codex", "app-server"]), + ], +) +@pytest.mark.asyncio +async def test_reply_rejects_call_time_client_options(tmp_path, name, value): + wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") + + with pytest.raises(TypeError, match=f"configured on the wrapper: {name}"): + await wrapper.reply("answer", **{name: value}) + + assert wrapper._codex is None + + +def test_constructor_rejects_launch_args_override(tmp_path): + with pytest.raises(TypeError, match="configure codex_bin instead"): + _wrapper(tmp_path, launch_args_override=["codex", "app-server"]) + def test_oauth_mode_ignores_api_credentials_and_forces_chatgpt(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path) + wrapper, _job = _wrapper( + tmp_path, + auth_mode="oauth", + api_key="explicit-key", + base_url="https://explicit.example.test/v1", + ) + wrapper.app_context.app_config.environment = {"TOOL_ENV": "preserved"} monkeypatch.setenv("CODEX_API_KEY", "ambient-key") monkeypatch.setenv("CODEX_BASE_URL", "https://ambient.example.test/v1") - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.load_env", lambda *_args: {}) auth = wrapper._resolve_auth_config( # pylint: disable=protected-access - { - "auth_mode": "oauth", - "api_key": "explicit-key", - "base_url": "https://explicit.example.test/v1", - }, + "oauth", + "explicit-key", + "https://explicit.example.test/v1", ) - config = wrapper._build_client_config({}, auth) # pylint: disable=protected-access + config = wrapper._build_client_config(auth) # pylint: disable=protected-access assert auth.mode == "oauth" assert auth.api_key == "" assert auth.base_url == "" - assert "OPENAI_API_KEY" not in config.env + assert config.env["TOOL_ENV"] == "preserved" + assert "CODEX_HOME" not in wrapper.app_context.app_config.environment assert 'forced_login_method="chatgpt"' in config.config_overrides assert not any(value.startswith("openai_base_url=") for value in config.config_overrides) @pytest.mark.asyncio async def test_api_key_mode_logs_in_app_server_explicitly(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path) + wrapper, _job = _wrapper( + tmp_path, + auth_mode="api_key", + api_key="explicit-key", + base_url="https://proxy.example.test/v1", + ) observed = {} class FakeCodex: @@ -718,34 +926,42 @@ async def test_api_key_mode_logs_in_app_server_explicitly(tmp_path, monkeypatch) async def close(self): observed["closed"] = True - monkeypatch.setattr("openai_codex.AsyncCodex", FakeCodex) - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.load_env", lambda *_args: {}) + monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) await wrapper.start() - await wrapper._get_codex( # pylint: disable=protected-access - { - "auth_mode": "api_key", - "api_key": "explicit-key", - "base_url": "https://proxy.example.test/v1", - }, - ) + await wrapper._get_codex() # pylint: disable=protected-access await wrapper.close() config = observed["config"] assert observed["api_key"] == "explicit-key" - assert "OPENAI_API_KEY" not in config.env assert 'openai_base_url="https://proxy.example.test/v1"' in config.config_overrides assert 'forced_login_method="api"' in config.config_overrides assert observed["closed"] is True -def test_api_key_mode_requires_key(tmp_path, monkeypatch): +def test_auth_selection_uses_explicit_wrapper_options(tmp_path, monkeypatch): wrapper, _job = _wrapper(tmp_path) + wrapper.app_context.app_config.components[ComponentEnum.AS_LLM]["default"] = SimpleNamespace( + credential={"api_key": "default-key", "base_url": "https://default.example.test/v1"}, + ) for name in ("CODEX_API_KEY", "OPENAI_API_KEY", "LLM_API_KEY"): - monkeypatch.delenv(name, raising=False) + monkeypatch.setenv(name, "ambient-key") + for name in ("CODEX_BASE_URL", "OPENAI_BASE_URL", "LLM_BASE_URL"): + monkeypatch.setenv(name, "https://ambient.example.test/v1") + + auth = wrapper._resolve_auth_config( # pylint: disable=protected-access + "api_key", + "configured-key", + ) + + assert auth.mode == "api_key" + assert auth.api_key == "configured-key" + assert auth.base_url == "" with pytest.raises(ValueError, match="requires a non-empty API key"): - wrapper._resolve_auth_config({"auth_mode": "api_key"}) # pylint: disable=protected-access + wrapper._resolve_auth_config( # pylint: disable=protected-access + "api_key", + ) @pytest.mark.parametrize("review_status", ["approved", "denied"]) diff --git a/tests/unit/test_component_registry.py b/tests/unit/test_component_registry.py index 5a224314..d7cda5d1 100644 --- a/tests/unit/test_component_registry.py +++ b/tests/unit/test_component_registry.py @@ -50,13 +50,6 @@ def test_register_decorator(): assert reg.get(ComponentEnum.FILE_CHUNKER, "alias") is MyParser -def test_register_overwrite_warns(caplog): - reg = ComponentRegistry() - reg.register(_DummyComponent, "dup") - reg.register(_DummyComponent, "dup") - assert reg.get(ComponentEnum.FILE_CHUNKER, "dup") is _DummyComponent - - def test_register_rejects_missing_component_type(): reg = ComponentRegistry() with pytest.raises(TypeError, match="ComponentEnum"): diff --git a/tests/unit/test_cron_job.py b/tests/unit/test_cron_job.py deleted file mode 100644 index e4be4f5a..00000000 --- a/tests/unit/test_cron_job.py +++ /dev/null @@ -1,77 +0,0 @@ -"""Unit tests for the ``cron`` job.""" - -# pylint: disable=protected-access - -import asyncio -from types import SimpleNamespace - -from reme.components import R -from reme.components.job.base_job import BaseJob -from reme.components.job.cron_job import CronJob -from reme.steps.base_step import BaseStep - - -@R.register("test_cron_counter_step") -class _CounterStep(BaseStep): - fires: int = 0 - - async def execute(self): - type(self).fires += 1 - self.context.response.success = True - return self.context.response - - -def _test_cron_parameter_protocol() -> None: - assert CronJob("* * * * *").cron_expr == "* * * * *" - assert CronJob(cron="0 3 * * *").cron_expr == "0 3 * * *" - print("OK cron_parameter_protocol") - - -async def _invalid_cron_raises() -> None: - try: - await CronJob(cron="not a cron")._start() - except ValueError: - return - raise AssertionError("expected ValueError for invalid cron expression") - - -def _test_invalid_cron_raises_on_start() -> None: - asyncio.run(_invalid_cron_raises()) - print("OK invalid_cron_raises_on_start") - - -async def _drive_steps_once() -> int: - _CounterStep.fires = 0 - job = CronJob( - cron="* * * * *", - steps=[{"backend": "test_cron_counter_step"}], - ) - job.app_context = SimpleNamespace(app_config=SimpleNamespace(language="")) - await BaseJob._start(job) - job._stop_event = asyncio.Event() - job._next_fire_delay = lambda: 0.01 - - task = asyncio.create_task(job()) - await asyncio.sleep(0.05) - job._stop_event.set() - await asyncio.wait_for(task, timeout=1) - return _CounterStep.fires - - -def _test_executes_own_steps() -> None: - count = asyncio.run(_drive_steps_once()) - assert count >= 1, f"expected cron to execute own steps, got {count}" - print(f"OK executes_own_steps count={count}") - - -def main() -> None: - """Run all tests.""" - print("=== cron job unit tests ===") - _test_cron_parameter_protocol() - _test_invalid_cron_raises_on_start() - _test_executes_own_steps() - print("=== passed ===") - - -if __name__ == "__main__": - main() diff --git a/tests/unit/test_default_file_chunker.py b/tests/unit/test_default_file_chunker.py index 31b608d1..1c300048 100644 --- a/tests/unit/test_default_file_chunker.py +++ b/tests/unit/test_default_file_chunker.py @@ -52,29 +52,6 @@ def test_parse_small_file(): asyncio.run(run()) -def test_parse_multiline_file(): - """Test parsing a file with multiple lines.""" - - async def run(): - lines = ["Line 1", "Line 2", "Line 3", "Line 4", "Line 5"] - content = "\n".join(lines) - with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f: - f.write(content) - temp_path = f.name - - try: - chunker = DefaultFileChunker(chunk_byte_size=10000) - _, chunks = await chunker.chunk(temp_path) - assert len(chunks) == 1 - assert chunks[0].start_line == 1 - assert chunks[0].end_line == 5 - print("✓ test_parse_multiline_file passed") - finally: - os.unlink(temp_path) - - asyncio.run(run()) - - def test_parse_chunked_file(): """Test parsing a file that requires multiple chunks.""" @@ -120,55 +97,6 @@ def test_parse_with_custom_encoding(): asyncio.run(run()) -def test_file_node_properties(): - """Test FileNode has correct properties.""" - - async def run(): - content = "test content" - with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f: - f.write(content) - temp_path = f.name - - try: - chunker = DefaultFileChunker() - file_node, _ = await chunker.chunk(temp_path) - assert hasattr(file_node, "path") - assert hasattr(file_node, "st_mtime") - assert file_node.st_mtime > 0 - print("✓ test_file_node_properties passed") - finally: - os.unlink(temp_path) - - asyncio.run(run()) - - -def test_file_chunk_properties(): - """Test FileChunk has correct properties.""" - - async def run(): - content = "test content for chunk" - with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f: - f.write(content) - temp_path = f.name - - try: - chunker = DefaultFileChunker() - _, chunks = await chunker.chunk(temp_path) - chunk = chunks[0] - assert hasattr(chunk, "path") - assert hasattr(chunk, "start_line") - assert hasattr(chunk, "end_line") - assert hasattr(chunk, "text") - assert hasattr(chunk, "id") - assert chunk.start_line >= 1 - assert chunk.end_line >= chunk.start_line - print("✓ test_file_chunk_properties passed") - finally: - os.unlink(temp_path) - - asyncio.run(run()) - - def test_parse_links_bare(): """Bare wikilink: [[target]].""" links = WikilinkHandler.extract_links("see [[note]]", "src.md") @@ -450,11 +378,8 @@ def test_min_chunk_and_overlap_size(): if __name__ == "__main__": test_parse_empty_file() test_parse_small_file() - test_parse_multiline_file() test_parse_chunked_file() test_parse_with_custom_encoding() - test_file_node_properties() - test_file_chunk_properties() test_parse_links_bare() test_parse_links_with_anchor() test_parse_links_alias_dropped() diff --git a/tests/unit/test_embedded_consumer_compat.py b/tests/unit/test_embedded_consumer_compat.py new file mode 100644 index 00000000..2570d04c --- /dev/null +++ b/tests/unit/test_embedded_consumer_compat.py @@ -0,0 +1,80 @@ +"""Compatibility tests for applications that embed ReMe in-process.""" + +import asyncio + +from reme import ReMe +from reme.components.agent_wrapper import AsAgentWrapper +from reme.enumeration import ComponentEnum + + +def _qwenpaw_style_config(workspace_dir: str) -> dict: + """Return the narrow ReMe contract used by QwenPaw's memory manager.""" + return { + "workspace_dir": workspace_dir, + "enable_logo": False, + "log_to_console": False, + "log_to_file": False, + "service": {"backend": "http"}, + "jobs": { + "version": { + "backend": "base", + "description": "return reme package version", + "parameters": {"type": "object", "properties": {}}, + "steps": [{"backend": "version_step"}], + }, + }, + "components": { + "as_llm": { + "default": { + "backend": "openai", + "model": "consumer-injected", + "credential": {"api_key": "", "base_url": ""}, + }, + }, + "agent_wrapper": { + "default": { + "backend": "agentscope", + "as_llm": "default", + }, + }, + }, + } + + +def test_qwenpaw_style_config_preserves_optional_defaults(tmp_path): + """New application fields remain optional for existing embedded configs.""" + app = ReMe(**_qwenpaw_style_config(str(tmp_path))) + + assert app.config.environment == {} + assert app.context.service is not None + assert app.context.service.jobs is None + + wrapper = app.context.components[ComponentEnum.AGENT_WRAPPER]["default"] + assert isinstance(wrapper, AsAgentWrapper) + assert wrapper.subprocess_environment == {} + + +def test_qwenpaw_style_config_keeps_in_process_application_api(tmp_path): + """Model injection, lifecycle, and direct job execution remain compatible.""" + app = ReMe(**_qwenpaw_style_config(str(tmp_path))) + injected_model = object() + + async def exercise_api() -> None: + component = await app.update_component( + "as_llm", + "default", + model=injected_model, + ) + await app.start() + try: + response = await app.run_job("version") + + assert component.model is injected_model + assert response.success is True + assert response.answer + finally: + await app.close() + + assert app.is_started is False + + asyncio.run(exercise_api()) diff --git a/tests/unit/test_hermes_agent_plugin.py b/tests/unit/test_hermes_agent_plugin.py deleted file mode 100644 index d9507796..00000000 --- a/tests/unit/test_hermes_agent_plugin.py +++ /dev/null @@ -1,457 +0,0 @@ -"""Hermes provider contract tests using a real loopback HTTP server.""" - -# pylint: disable=redefined-outer-name - -from __future__ import annotations - -import importlib.util -import http.client -import json -import sys -import threading -import time -import types - -from collections.abc import Iterator -from contextlib import contextmanager -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path -from typing import Any - -import pytest - - -class _MemoryProvider: - """Minimal Hermes ABC stand-in; the real loader is exercised separately.""" - - -@pytest.fixture -def plugin_module(monkeypatch: pytest.MonkeyPatch): - """Load the plugin the same way Hermes loads an isolated plugin package.""" - agent = types.ModuleType("agent") - agent.__path__ = [] # type: ignore[attr-defined] - memory_provider = types.ModuleType("agent.memory_provider") - memory_provider.MemoryProvider = _MemoryProvider - monkeypatch.setitem(sys.modules, "agent", agent) - monkeypatch.setitem(sys.modules, "agent.memory_provider", memory_provider) - - module_name = "_reme_hermes_test_plugin" - for name in list(sys.modules): - if name == module_name or name.startswith(f"{module_name}."): - monkeypatch.delitem(sys.modules, name, raising=False) - plugin_dir = Path(__file__).resolve().parents[2] / "plugins" / "hermes_agent" - spec = importlib.util.spec_from_file_location( - module_name, - plugin_dir / "__init__.py", - submodule_search_locations=[str(plugin_dir)], - ) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, module_name, module) - spec.loader.exec_module(module) - return module - - -class _ActionHandler(BaseHTTPRequestHandler): - calls: list[tuple[str, dict[str, Any]]] - responses: dict[str, tuple[int, Any] | tuple[int, Any, float]] - - def do_POST(self) -> None: # noqa: N802 - stdlib callback name - """Serve one ReMe-compatible action request.""" - length = int(self.headers.get("Content-Length", "0")) - body = json.loads(self.rfile.read(length) or b"{}") - self.calls.append((self.path, body)) - spec = self.responses.get(self.path, (404, {"detail": "not found"})) - status, response = spec[:2] - if len(spec) == 3: - time.sleep(spec[2]) - encoded = json.dumps(response).encode("utf-8") - self.send_response(status) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(encoded))) - self.end_headers() - try: - self.wfile.write(encoded) - except (BrokenPipeError, ConnectionResetError): - pass - - def log_message(self, _format: str, *_args: object) -> None: - return - - -@contextmanager -def action_server( - responses: dict[str, tuple[int, Any] | tuple[int, Any, float]], -) -> Iterator[tuple[str, list[tuple[str, dict[str, Any]]]]]: - """Run a loopback ReMe action server and expose captured requests.""" - calls: list[tuple[str, dict[str, Any]]] = [] - handler = type("Handler", (_ActionHandler,), {"calls": calls, "responses": responses}) - server = ThreadingHTTPServer(("127.0.0.1", 0), handler) - thread = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - try: - host, port = server.server_address - yield f"http://{host}:{port}", calls - finally: - server.shutdown() - server.server_close() - thread.join(timeout=2) - - -def _healthy_response(answer: str = "healthy") -> dict[str, Any]: - return { - "success": True, - "answer": answer, - "metadata": {"health": {"healthy": True, "version": "test"}}, - } - - -def _wait_for_call(calls: list[tuple[str, dict[str, Any]]], path: str, timeout: float = 2.0) -> None: - """Wait until the background writer reaches a loopback action.""" - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if any(call_path == path for call_path, _ in calls): - return - time.sleep(0.01) - raise AssertionError(f"timed out waiting for {path}; calls={calls!r}") - - -def _write_config(plugin_module, home: Path, endpoint: str, **overrides: Any) -> None: - """Seed runtime config without exercising the separately tested setup path.""" - del plugin_module - values = {"endpoint": endpoint, "health_retry_seconds": 0.1, **overrides} - home.mkdir(parents=True, exist_ok=True) - (home / "reme.json").write_text(json.dumps(values), encoding="utf-8") - - -def test_setup_is_profile_scoped_and_atomic(plugin_module, tmp_path: Path) -> None: - """Keep endpoint configuration private to each Hermes profile.""" - first = tmp_path / "profile-a" - second = tmp_path / "profile-b" - provider = plugin_module.ReMeMemoryProvider() - - responses = {"/health_check": (200, _healthy_response())} - with action_server(responses) as (first_endpoint, _), action_server(responses) as (second_endpoint, _): - provider.save_config({"endpoint": first_endpoint}, str(first)) - provider.save_config({"endpoint": second_endpoint}, str(second)) - - first_config = json.loads((first / "reme.json").read_text(encoding="utf-8")) - second_config = json.loads((second / "reme.json").read_text(encoding="utf-8")) - assert first_config["endpoint"] == first_endpoint - assert second_config["endpoint"] == second_endpoint - assert (first / "reme.json").stat().st_mode & 0o777 == 0o600 - assert not list(first.glob(".reme.json.*")) - - -def test_setup_rejects_unhealthy_endpoint_without_overwrite(plugin_module, tmp_path: Path) -> None: - """Preserve the last working profile config when endpoint validation fails.""" - healthy = {"/health_check": (200, _healthy_response())} - unhealthy = {"/health_check": (503, {"detail": "starting"})} - provider = plugin_module.ReMeMemoryProvider() - - with action_server(healthy) as (healthy_endpoint, _), action_server(unhealthy) as (unhealthy_endpoint, _): - provider.save_config({"endpoint": healthy_endpoint}, str(tmp_path)) - with pytest.raises(plugin_module.ReMeServiceError, match="HTTP 503"): - provider.save_config({"endpoint": unhealthy_endpoint}, str(tmp_path)) - - saved = json.loads((tmp_path / "reme.json").read_text(encoding="utf-8")) - assert saved["endpoint"] == healthy_endpoint - - -def test_setup_rejects_incomplete_health_envelope(plugin_module, tmp_path: Path) -> None: - """Do not accept an unrelated endpoint that returns generic JSON.""" - responses = {"/health_check": (200, {"success": True, "answer": "ok"})} - with action_server(responses) as (endpoint, _): - with pytest.raises(plugin_module.ReMeServiceError, match="healthy component snapshot"): - plugin_module.ReMeMemoryProvider().save_config({"endpoint": endpoint}, str(tmp_path)) - - -def test_client_requires_explicit_success_envelope(plugin_module) -> None: - """Reject action responses that omit ReMe's explicit success signal.""" - responses = {"/search": (200, {"answer": "not a ReMe envelope"})} - with action_server(responses) as (endpoint, _): - client = plugin_module.ReMeHttpClient(endpoint, timeout=1.0) - with pytest.raises(plugin_module.ReMeServiceError, match="not a ReMe envelope"): - client.call("search", {"query": "test"}) - - -def test_lifecycle_retrieves_records_and_switches_sessions(plugin_module, tmp_path: Path) -> None: - """Exercise recall, recording, session switching, and shutdown.""" - responses = { - "/health_check": (200, _healthy_response()), - "/search": (200, {"success": True, "answer": "remembered project decision", "metadata": {}}), - "/auto_memory": (200, {"success": True, "answer": "recorded", "metadata": {}}), - } - with action_server(responses) as (endpoint, calls): - _write_config(plugin_module, tmp_path, endpoint) - provider = plugin_module.ReMeMemoryProvider() - provider.initialize("conversation-one", hermes_home=str(tmp_path), agent_identity="coder") - - assert provider.prefetch("What did we decide?") == "remembered project decision" - provider.sync_turn("Use SQLite", "Recorded that decision") - provider.on_session_switch("conversation-two") - provider.sync_turn("Use BM25 too", "Recorded the retrieval choice") - provider.shutdown() - - assert [path for path, _ in calls] == [ - "/health_check", - "/search", - "/auto_memory", - "/auto_memory", - ] - first_record = calls[2][1] - second_record = calls[3][1] - assert first_record["session_id"].startswith("hermes-coder-conversation-one-") - assert second_record["session_id"].startswith("hermes-coder-conversation-two-") - assert first_record["session_id"] != second_record["session_id"] - assert first_record["messages"] == [ - {"name": "user", "role": "user", "content": "Use SQLite"}, - {"name": "assistant", "role": "assistant", "content": "Recorded that decision"}, - ] - - -def test_gateway_session_argument_preserves_conversation_boundary(plugin_module, tmp_path: Path) -> None: - """Use per-request gateway sessions instead of cached provider state.""" - responses = { - "/health_check": (200, _healthy_response()), - "/auto_memory": (200, {"success": True, "answer": "recorded", "metadata": {}}), - } - with action_server(responses) as (endpoint, calls): - _write_config(plugin_module, tmp_path, endpoint) - provider = plugin_module.ReMeMemoryProvider() - provider.initialize("cached-agent", hermes_home=str(tmp_path), agent_identity="gateway") - provider.sync_turn("first", "one", session_id="chat-a") - provider.sync_turn("second", "two", session_id="chat-b") - provider.shutdown() - - assert calls[1][1]["session_id"].startswith("hermes-gateway-chat-a-") - assert calls[2][1]["session_id"].startswith("hermes-gateway-chat-b-") - - -def test_non_primary_context_does_not_write(plugin_module, tmp_path: Path) -> None: - """Avoid recording internal cron, flush, and subagent turns.""" - responses = { - "/health_check": (200, _healthy_response()), - "/auto_memory": (200, {"success": True, "answer": "recorded", "metadata": {}}), - } - with action_server(responses) as (endpoint, calls): - _write_config(plugin_module, tmp_path, endpoint) - provider = plugin_module.ReMeMemoryProvider() - provider.initialize( - "cron-session", - hermes_home=str(tmp_path), - agent_identity="default", - agent_context="cron", - ) - provider.sync_turn("scheduled system prompt", "scheduled result") - - assert [path for path, _ in calls] == ["/health_check"] - - -def test_unavailable_service_fails_open_and_reports_dropped_write(plugin_module, tmp_path: Path, caplog) -> None: - """Keep Hermes usable while making lost persistence explicit.""" - responses = {"/health_check": (503, {"detail": "starting"})} - with action_server(responses) as (endpoint, calls): - _write_config(plugin_module, tmp_path, endpoint) - provider = plugin_module.ReMeMemoryProvider() - with caplog.at_level("WARNING"): - provider.initialize("session", hermes_home=str(tmp_path), agent_identity="default") - assert provider.prefetch("question") == "" - provider.sync_turn("important fact", "answer") - provider.shutdown() - - assert [path for path, _ in calls] == ["/health_check"] - assert "recall is disabled" in caplog.text - assert "did not record completed turn" in caplog.text - - -def test_unhealthy_snapshot_is_not_treated_as_available(plugin_module, tmp_path: Path) -> None: - """Reject a successful HTTP response whose health snapshot is unhealthy.""" - responses = { - "/health_check": ( - 200, - { - "success": True, - "answer": "unhealthy", - "metadata": {"health": {"healthy": False}}, - }, - ), - } - with action_server(responses) as (endpoint, calls): - _write_config(plugin_module, tmp_path, endpoint) - provider = plugin_module.ReMeMemoryProvider() - provider.initialize("session", hermes_home=str(tmp_path), agent_identity="default") - assert provider.prefetch("question") == "" - - assert [path for path, _ in calls] == ["/health_check"] - - -def test_recall_timeout_does_not_block_hermes_turn(plugin_module, tmp_path: Path) -> None: - """Bound inline recall independently from slow automatic-memory writes.""" - responses = { - "/health_check": (200, _healthy_response()), - "/search": (200, {"success": True, "answer": "too late", "metadata": {}}, 0.5), - "/auto_memory": (200, {"success": True, "answer": "recorded", "metadata": {}}), - } - with action_server(responses) as (endpoint, calls): - _write_config(plugin_module, tmp_path, endpoint, recall_timeout=0.1) - provider = plugin_module.ReMeMemoryProvider() - provider.initialize("session", hermes_home=str(tmp_path), agent_identity="default") - - started = time.monotonic() - assert provider.prefetch("question") == "" - elapsed = time.monotonic() - started - assert provider.prefetch("cooldown") == "" - provider.sync_turn("recall timed out", "write still works") - provider.shutdown() - - assert elapsed < 0.4 - assert [path for path, _ in calls] == ["/health_check", "/search", "/auto_memory"] - - -def test_recording_failure_does_not_disable_recall(plugin_module, tmp_path: Path) -> None: - """Keep retrieval healthy when only ReMe's LLM-backed write path fails.""" - responses = { - "/health_check": (200, _healthy_response()), - "/auto_memory": (200, {"success": False, "answer": "LLM unavailable", "metadata": {}}), - "/search": (200, {"success": True, "answer": "recall still works", "metadata": {}}), - } - with action_server(responses) as (endpoint, calls): - _write_config(plugin_module, tmp_path, endpoint) - provider = plugin_module.ReMeMemoryProvider() - provider.initialize("session", hermes_home=str(tmp_path), agent_identity="default") - provider.sync_turn("remember this", "attempted") - _wait_for_call(calls, "/auto_memory") - - assert provider.prefetch("existing fact") == "recall still works" - provider.shutdown() - - assert [path for path, _ in calls] == ["/health_check", "/auto_memory", "/search"] - - -def test_writer_continues_after_unexpected_payload_exception( - plugin_module, - tmp_path: Path, - caplog, - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Do not let a malformed HTTP response kill all later writes.""" - responses = { - "/health_check": (200, _healthy_response()), - "/auto_memory": (200, {"success": True, "answer": "recorded", "metadata": {}}), - } - with action_server(responses) as (endpoint, calls): - _write_config(plugin_module, tmp_path, endpoint) - provider = plugin_module.ReMeMemoryProvider() - provider.initialize("session", hermes_home=str(tmp_path), agent_identity="default") - original_record = getattr(provider, "_record_payload") - attempts = 0 - - def record_with_one_broken_response(payload: dict[str, Any]) -> None: - nonlocal attempts - attempts += 1 - if attempts == 1: - raise http.client.IncompleteRead(b"partial", 10) - original_record(payload) - - monkeypatch.setattr(provider, "_record_payload", record_with_one_broken_response) - with caplog.at_level("ERROR"): - provider.sync_turn("first", "broken response") - provider.sync_turn("second", "must still be recorded") - provider.shutdown() - - assert attempts == 2 - assert "Unexpected ReMe recording failure" in caplog.text - assert [path for path, _ in calls] == ["/health_check", "/auto_memory"] - - -def test_enqueue_restarts_a_dead_writer(plugin_module, tmp_path: Path) -> None: - """Replace a stale worker reference before accepting another payload.""" - responses = { - "/health_check": (200, _healthy_response()), - "/auto_memory": (200, {"success": True, "answer": "recorded", "metadata": {}}), - } - with action_server(responses) as (endpoint, calls): - _write_config(plugin_module, tmp_path, endpoint) - provider = plugin_module.ReMeMemoryProvider() - provider.initialize("session", hermes_home=str(tmp_path), agent_identity="default") - dead_writer = threading.Thread(target=lambda: None) - dead_writer.start() - dead_writer.join(timeout=1) - assert not dead_writer.is_alive() - setattr(provider, "_write_thread", dead_writer) - - provider.sync_turn("after failure", "record this") - provider.shutdown() - - assert getattr(provider, "_write_thread") is None - assert [path for path, _ in calls] == ["/health_check", "/auto_memory"] - - -def test_shutdown_is_bounded_when_write_is_slow(plugin_module, tmp_path: Path, caplog) -> None: - """Do not let an in-flight automatic-memory request wedge Hermes exit.""" - responses = { - "/health_check": (200, _healthy_response()), - "/auto_memory": (200, {"success": True, "answer": "recorded", "metadata": {}}, 0.8), - } - with action_server(responses) as (endpoint, calls): - _write_config(plugin_module, tmp_path, endpoint, shutdown_timeout=0.1) - provider = plugin_module.ReMeMemoryProvider() - provider.initialize("session", hermes_home=str(tmp_path), agent_identity="default") - provider.sync_turn("slow", "write") - _wait_for_call(calls, "/auto_memory") - provider.sync_turn("queued", "must be abandoned") - - with caplog.at_level("WARNING"): - started = time.monotonic() - provider.shutdown() - elapsed = time.monotonic() - started - - assert elapsed < 0.4 - assert "abandoned 1 queued write(s)" in caplog.text - assert [path for path, _ in calls] == ["/health_check", "/auto_memory"] - - -def test_process_exit_fallback_drains_once(plugin_module, tmp_path: Path) -> None: - """Drain a queued turn when Hermes exits without normal provider cleanup.""" - responses = { - "/health_check": (200, _healthy_response()), - "/auto_memory": (200, {"success": True, "answer": "recorded", "metadata": {}}), - } - with action_server(responses) as (endpoint, calls): - _write_config(plugin_module, tmp_path, endpoint) - provider = plugin_module.ReMeMemoryProvider() - provider.initialize("session", hermes_home=str(tmp_path), agent_identity="default") - provider.sync_turn("process", "exit") - - shutdown_at_exit = getattr(provider, "_atexit_shutdown") - shutdown_at_exit() - shutdown_at_exit() # idempotent if normal cleanup also ran - - assert [path for path, _ in calls] == ["/health_check", "/auto_memory"] - - -def test_service_recovers_after_health_retry_cooldown(plugin_module, tmp_path: Path) -> None: - """Resume recall after a previously unavailable ReMe service recovers.""" - responses = {"/health_check": (503, {"detail": "starting"})} - with action_server(responses) as (endpoint, calls): - _write_config(plugin_module, tmp_path, endpoint) - provider = plugin_module.ReMeMemoryProvider() - provider.initialize("session", hermes_home=str(tmp_path), agent_identity="default") - - responses["/health_check"] = (200, _healthy_response()) - responses["/search"] = (200, {"success": True, "answer": "recovered", "metadata": {}}) - time.sleep(0.11) - assert provider.prefetch("question") == "recovered" - - assert [path for path, _ in calls] == ["/health_check", "/health_check", "/search"] - - -def test_register_exposes_provider(plugin_module) -> None: - """Register exactly one provider without model-visible tools.""" - registered: list[Any] = [] - context = types.SimpleNamespace(register_memory_provider=registered.append) - plugin_module.register(context) - assert len(registered) == 1 - assert registered[0].name == "reme" - assert registered[0].get_tool_schemas() == [] diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index fabb851a..06b13998 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -236,31 +236,6 @@ def test_supervisor_disabled_propagates_exception(): asyncio.run(run()) -# -- BackgroundJob._wait_or_stop ---------------------------------------------- - - -def test_wait_or_stop_returns_on_stop(): - async def run(): - job = BackgroundJob(name="bg") - job._stop_event = asyncio.Event() - job._stop_event.set() - await job._wait_or_stop(10.0) - - asyncio.run(run()) - - -# -- BackgroundJob._shutdown_task --------------------------------------------- - - -def test_shutdown_task_none(): - async def run(): - job = BackgroundJob(name="bg") - job._task = None - await job._shutdown_task() - - asyncio.run(run()) - - def test_shutdown_task_cancels_on_timeout(): async def run(): async def hang_forever(): @@ -357,22 +332,3 @@ def test_application_start_failure_propagates_and_closes_started_components(): assert not app._started_components asyncio.run(run()) - - -if __name__ == "__main__": - print("\n=== Job Tests ===") - test_resolve_step_missing_backend() - test_resolve_step_unregistered_backend() - test_call_captures_exception() - test_call_runs_steps_in_order() - test_start_without_app_context_raises() - test_backoff_delay_increases() - test_backoff_delay_capped() - test_backoff_delay_has_jitter() - test_backoff_delay_attempt_zero() - test_supervisor_restarts_on_crash() - test_supervisor_disabled_propagates_exception() - test_wait_or_stop_returns_on_stop() - test_shutdown_task_none() - test_shutdown_task_cancels_on_timeout() - print("\n所有测试通过!") diff --git a/tests/unit/test_json_file_chunker.py b/tests/unit/test_json_file_chunker.py index f32eca31..4d8af06a 100644 --- a/tests/unit/test_json_file_chunker.py +++ b/tests/unit/test_json_file_chunker.py @@ -474,33 +474,6 @@ class TestDfsAlgorithm: assert "a" in cdata assert isinstance(cdata["a"], list) - def test_dfs_order_no_skip(self): - """DFS order: cannot have a[0] and root key 'b' in same chunk - while skipping a[1]-a[3].""" - data = {"a": ["x" * 200, {"c": 1, "d": 2}, "y" * 200, 1], "b": 10} - chunker = JsonFileChunker(chunk_chars=300) - chunker.min_element_size = 1 - root = _build_tree(chunker, data) - chunks = chunker._node_to_chunks(root) - # No chunk should contain key "b" while only having part of "a" - for cdata, _, _ in chunks: - if "b" in cdata: - # If "b" is present, all "a" leaves must be in earlier chunks - # (this chunk must be the last or "a" is complete here) - a_val = cdata.get("a") - if isinstance(a_val, list): - # "a" is present — verify it's either complete or this - # is a continuation from previous chunk - pass # structural check only - # At minimum, verify no single chunk has ONLY a[0] and b:10 - for cdata, _, _ in chunks: - if "b" in cdata and "a" in cdata: - a_list = cdata["a"] - # If a has only one element and it's "x"*200 (a[0]), - # that would mean skipping a[1]-a[3] — not allowed - if len(a_list) == 1 and isinstance(a_list[0], str): - pytest.fail("DFS order violated: a[0] and b in same chunk without a[1]-a[3]") - def test_calibration_during_chunking(self): """Calibration checkpoint keeps size accurate for large inputs.""" data = {f"k{i}": "v" * 50 for i in range(200)} diff --git a/tests/unit/test_keyword_index.py b/tests/unit/test_keyword_index.py index 387a643d..607a9e36 100644 --- a/tests/unit/test_keyword_index.py +++ b/tests/unit/test_keyword_index.py @@ -68,38 +68,6 @@ def run(coro): # --------------------------------------------------------------------------- # -def test_basic_init(): - """Default constructor produces empty BM25 state.""" - - async def go(): - with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - bm25 = BM25Index() - assert bm25.k1 == 1.5 - assert bm25.b == 0.75 - assert bm25.index_version == "v1" - assert bm25.vocab == {} - assert not bm25.inverted_index - assert bm25.doc_meta == {} - assert bm25.n_docs == 0 - assert bm25.total_len == 0 - assert bm25.avg_len == 0.0 - - run(go()) - - -def test_custom_params(): - """k1, b and index_version are honoured.""" - - async def go(): - with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - bm25 = BM25Index(k1=2.0, b=0.5, index_version="v2") - assert bm25.k1 == 2.0 - assert bm25.b == 0.5 - assert bm25.index_version == "v2" - - run(go()) - - def test_index_file_raises_when_tokenizer_is_none(): """index_file must raise when tokenizer is explicitly None.""" diff --git a/tests/unit/test_logging_config.py b/tests/unit/test_logging_config.py index fe3a82ee..1293a74b 100644 --- a/tests/unit/test_logging_config.py +++ b/tests/unit/test_logging_config.py @@ -1,10 +1,12 @@ """Tests for logging configuration handoff during app startup.""" import concurrent.futures +from datetime import datetime import io import logging import threading import time +from unittest.mock import Mock import pytest @@ -25,6 +27,20 @@ class DummyLogger: return None +def test_loguru_filename_includes_start_time_and_process_id(monkeypatch, tmp_path): + """Independent ReMe processes should write to distinct Loguru files.""" + fixed_datetime = Mock() + fixed_datetime.now.return_value = datetime(2026, 7, 20, 15, 42, 18) + monkeypatch.setattr(logger_utils, "datetime", fixed_datetime) + monkeypatch.setattr(logger_utils.os, "getpid", lambda: 31247) + + logger = logger_utils._init_loguru(str(tmp_path), "INFO", False, True) # pylint: disable=protected-access + try: + assert (tmp_path / "2026-07-20_15-42-18_31247.log").is_file() + finally: + logger.remove() + + def test_stdlib_formatter_matches_qwenpaw_console_format(monkeypatch, tmp_path, capsys): """Stdlib logs should use QwenPaw's level/path/time/message layout.""" monkeypatch.chdir(tmp_path) diff --git a/tests/unit/test_reme_cli.py b/tests/unit/test_reme_cli.py index 03f3dd1a..95930f93 100644 --- a/tests/unit/test_reme_cli.py +++ b/tests/unit/test_reme_cli.py @@ -1,6 +1,9 @@ """Tests for the ReMe CLI entry helpers.""" import asyncio +from pathlib import Path +import subprocess +import sys from types import SimpleNamespace import pytest @@ -10,6 +13,41 @@ from reme.components.service.cli_service import CliService from reme import reme as reme_module +def test_package_import_does_not_require_optional_agent_sdks(): + """The base package remains importable without Claude or Codex SDKs.""" + script = """ +import importlib.abc +import sys + + +class BlockOptionalAgentSDKs(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + blocked = ("claude_agent_sdk", "openai_codex") + if any(fullname == name or fullname.startswith(f"{name}.") for name in blocked): + raise ModuleNotFoundError(f"blocked optional SDK: {fullname}", name=fullname) + return None + + +sys.meta_path.insert(0, BlockOptionalAgentSDKs()) +import reme + +assert not any( + name == sdk or name.startswith(f"{sdk}.") + for name in sys.modules + for sdk in ("claude_agent_sdk", "openai_codex") +) +""" + result = subprocess.run( + [sys.executable, "-c", script], + cwd=Path(__file__).resolve().parents[2], + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, result.stderr + + def test_main_loads_env_before_calling_server(monkeypatch): """Client actions can resolve connection settings from the local .env.""" events = [] @@ -28,6 +66,38 @@ def test_main_loads_env_before_calling_server(monkeypatch): assert events == ["load_env", ("call_server", "shell", {"cmd": "pwd"})] +def test_main_saves_loaded_environment_in_start_config(monkeypatch): + """The startup config keeps the environment captured by the single global load.""" + observed = {} + + class FakeReMe: + """Capture the fully resolved application configuration.""" + + def __init__(self, **kwargs): + """Record the application startup configuration.""" + observed["config"] = kwargs + + def run_app(self): + """Record that application startup continued.""" + observed["ran"] = True + + main_globals = reme_module.main.__globals__ + monkeypatch.setitem(main_globals, "load_env", lambda: {"TOOL_ENV": "configured"}) + monkeypatch.setitem(main_globals, "parse_args", lambda *_args: ("start", {})) + monkeypatch.setitem(main_globals, "prepare_start_config", lambda _kwargs: {"service": {"backend": "cli"}}) + monkeypatch.setitem(main_globals, "ReMe", FakeReMe) + + reme_module.main() + + assert observed == { + "config": { + "service": {"backend": "cli"}, + "environment": {"TOOL_ENV": "configured"}, + }, + "ran": True, + } + + def test_prepare_start_config_moves_unknown_start_args_to_job_args(monkeypatch): """``reme start job=...`` is translated into a one-shot cli service config.""" diff --git a/tests/unit/test_service.py b/tests/unit/test_service.py index da8d564a..7dc4c294 100644 --- a/tests/unit/test_service.py +++ b/tests/unit/test_service.py @@ -1,9 +1,14 @@ """Tests for service job registration behavior.""" +import asyncio from types import SimpleNamespace +from unittest.mock import Mock + +import pytest from reme.components.job import BaseJob, StreamJob from reme.components.service import MCPService +from reme.schema import Response def _dummy_app(): @@ -23,6 +28,64 @@ def _dummy_app(): ) +def _app_with_jobs(**jobs): + """Minimal object needed by BaseService.add_jobs.""" + return SimpleNamespace(context=SimpleNamespace(jobs=jobs)) + + +def test_service_registers_all_enabled_jobs_by_default(): + """Omitting service.jobs preserves registration of every service-enabled job.""" + service = MCPService() + service.add_job = Mock(return_value=True) + enabled = BaseJob(name="enabled") + disabled = BaseJob(name="disabled", enable_serve=False) + + service.add_jobs(_app_with_jobs(enabled=enabled, disabled=disabled)) + + service.add_job.assert_called_once_with(enabled) + + +def test_service_jobs_restricts_registration_to_configured_names(): + """service.jobs acts as a whitelist without overriding enable_serve.""" + service = MCPService(jobs=["selected"]) + service.add_job = Mock(return_value=True) + selected = BaseJob(name="selected") + unselected = BaseJob(name="unselected") + disabled = BaseJob(name="disabled", enable_serve=False) + + service.add_jobs( + _app_with_jobs(selected=selected, unselected=unselected, disabled=disabled), + ) + + service.add_job.assert_called_once_with(selected) + + +def test_empty_service_jobs_disables_job_registration(): + """An explicit empty whitelist exposes no jobs.""" + service = MCPService(jobs=[]) + service.add_job = Mock(return_value=True) + + service.add_jobs(_app_with_jobs(enabled=BaseJob(name="enabled"))) + + service.add_job.assert_not_called() + + +def test_explicit_service_jobs_reject_missing_disabled_and_unsupported_jobs(): + """An explicit service.jobs list fails instead of starting an incomplete service.""" + missing_service = MCPService(jobs=["missing"]) + with pytest.raises(KeyError, match="missing"): + missing_service.add_jobs(_app_with_jobs()) + + disabled_service = MCPService(jobs=["disabled"]) + with pytest.raises(ValueError, match="disabled"): + disabled_service.add_jobs(_app_with_jobs(disabled=BaseJob(name="disabled", enable_serve=False))) + + stream_service = MCPService(jobs=["stream"]) + stream_service.add_job = Mock(return_value=False) + with pytest.raises(TypeError, match="stream"): + stream_service.add_jobs(_app_with_jobs(stream=StreamJob(name="stream"))) + + def test_mcp_service_registers_job_with_empty_parameters(): """Empty job parameters must remain a dict for FastMCP FunctionTool validation.""" service = MCPService() @@ -41,3 +104,89 @@ def test_mcp_service_reports_stream_job_skipped(): job = StreamJob(name="stream") assert service.add_job(job) is False + + +class _RecordingJob: + """Small callable matching the job contract used by MCPService.add_job.""" + + name = "record" + description = "Record arguments" + parameters = {"type": "object", "properties": {"query": {"type": "string"}}} + + def __init__(self, response: Response | None = None): + self.response = response or Response(answer="ok") + self.calls = [] + + async def __call__(self, **kwargs): + self.calls.append(kwargs) + return self.response + + +def test_mcp_service_injects_job_kwargs_and_rejects_conflicts(): + """Configured job arguments are injected exactly once and remain server-owned.""" + + async def run(): + service = MCPService(injected_job_kwargs={"tool_context_id": "ctx-1"}) + service.build_service(_dummy_app()) + job = _RecordingJob() + job.parameters = { + "type": "object", + "properties": { + "query": {"type": "string"}, + "tool_context_id": {"type": "string"}, + }, + "required": ["query", "tool_context_id"], + } + assert service.add_job(job) is True + tool = await service.service.get_tool(job.name) + assert tool is not None + assert "tool_context_id" not in tool.parameters["properties"] + assert tool.parameters["required"] == ["query"] + + result = await tool.run({"query": "alpha"}) + assert job.calls == [{"query": "alpha", "tool_context_id": "ctx-1"}] + assert "ok" in str(result.content) + + with pytest.raises(Exception, match="tool_context_id injected by the MCP server"): + await tool.run({"query": "alpha", "tool_context_id": "caller"}) + + asyncio.run(run()) + + +def test_mcp_service_can_raise_tool_error_for_unsuccessful_response(): + """Configured MCP services translate failed Responses into tool errors.""" + + async def run(): + service = MCPService(tool_error_on_failure=True) + service.build_service(_dummy_app()) + job = _RecordingJob(Response(answer="failed", success=False)) + assert service.add_job(job) is True + tool = await service.service.get_tool(job.name) + assert tool is not None + + with pytest.raises(Exception, match="failed"): + await tool.run({}) + + asyncio.run(run()) + + +def test_service_lifespan_closes_app_after_error(): + """Application resources close even when serving exits with an exception.""" + + async def run(): + events = [] + + async def start(): + events.append("start") + + async def close(): + events.append("close") + + app = SimpleNamespace(start=start, close=close) + lifespan = MCPService()._lifespan(app, "127.0.0.1", 0) # pylint: disable=protected-access + with pytest.raises(RuntimeError, match="stop"): + async with lifespan(None): + raise RuntimeError("stop") + assert events == ["start", "close"] + + asyncio.run(run())