ReMe/reme/components/agent_wrapper/cc_agent_wrapper.py
jinliyl c1a25e9ff4
feat(agent): add Codex agent wrapper and ReMe MCP bridge (#358)
* feat(agent): add Codex wrapper integration

* feat(agent): enhance agent wrapper functionality and add comprehensive testing

- Implement structured output schema normalization across all wrappers
- Add Claude Code system prompt mode support with append/replace options
- Introduce Codex agent wrapper with streaming, tool context isolation, and skill management
- Enhance skill linking with validation and conflict resolution
- Add approval event streaming support for Codex wrapper
- Implement output schema validation and normalize function
- Create dedicated test suites for Claude Code and Codex integration
- Update README documentation for Codex wrapper capabilities
- Refactor kwargs merging with proper schema handling
- Add tool context validation when resuming sessions
- Implement proper cleanup and session management for Codex wrapper

* test(cc-agent): add test coverage for structured output scenarios

- Add docstring for empty schema validation in build_options
- Document falsy structured output preservation behavior
- Add docstring for streaming wrapper schema rejection
- Include lambda function reference for wrapper factory consistency
- Add test documentation for live Codex wrapper contract exercise

* docs: revert README changes

* fix(agent): interrupt abandoned Codex turns
2026-07-17 13:39:18 +08:00

657 lines
28 KiB
Python

"""Claude Code SDK backend for the unified agent wrapper."""
import json
import os
from collections.abc import AsyncGenerator
from dataclasses import asdict
from pathlib import Path
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 ..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.
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
@R.register("claude_code")
class CcAgentWrapper(BaseAgentWrapper):
"""Agent wrapper backed by Claude Code 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,
}
@property
def session_path(self) -> Path:
"""Directory used for persisted Claude Code sessions."""
if self.app_context is None:
return self.workspace_path / "mem_session"
return self.workspace_path / self.app_context.app_config.mem_session_dir
def _ensure_claude_skill_dir(self, config_dir: Path, skills: list[str] | str) -> None:
"""Add selected project skills to Claude Code discovery locations."""
project_skills = self.project_skills_root
if not project_skills.is_dir():
return
if skills == "all":
skill_names = sorted(path.name for path in project_skills.iterdir() if path.is_dir())
else:
skill_names = list(dict.fromkeys(skills))
for skill_name in skill_names:
if not skill_name or Path(skill_name).name != skill_name or skill_name in {".", ".."}:
raise ValueError(f"Invalid skill name: {skill_name!r}")
sources = {
skill_name: project_skills / skill_name
for skill_name in skill_names
if (project_skills / skill_name).is_dir()
}
if not sources:
return
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 path: {target}")
continue
target.mkdir(parents=True, exist_ok=True)
for skill_name, source in sources.items():
skill_target = target / skill_name
if skill_target.is_symlink():
if skill_target.resolve() != source.resolve():
self.logger.warning(f"Preserving existing Claude Code skill link: {skill_target}")
continue
if skill_target.exists():
self.logger.warning(f"Preserving existing Claude Code skill path: {skill_target}")
continue
skill_target.symlink_to(source, target_is_directory=True)
except OSError as exc:
self.logger.warning(f"Failed to link Claude Code skills into {target}: {exc}")
@classmethod
def _make_tool(cls, job: "BaseJob", tool_context_id: str | None = None):
from claude_agent_sdk import SdkMcpTool
async def run_job(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}
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.
``stream=True`` enables ``include_partial_messages`` so that
``StreamEvent`` messages are emitted alongside the final
``ResultMessage``.
"""
from claude_agent_sdk import create_sdk_mcp_server
from claude_agent_sdk.types import ClaudeAgentOptions
self._apply_system_prompt_mode(kwargs)
skills = kwargs.get("skills")
if isinstance(skills, str) and skills != "all":
kwargs["skills"] = [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)
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:
extra_env_dict.update(
{
"ANTHROPIC_MODEL": model,
"ANTHROPIC_DEFAULT_HAIKU_MODEL": model,
"ANTHROPIC_DEFAULT_SONNET_MODEL": model,
"ANTHROPIC_DEFAULT_OPUS_MODEL": model,
},
)
opts.env.update(extra_env_dict)
self.session_path.mkdir(parents=True, exist_ok=True)
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")
job_tools: list[str] = kwargs.get("job_tools", [])
resolved_jobs = self._resolve_job_tools(job_tools)
if resolved_jobs:
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.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 -------------------------------------------
@classmethod
# pylint: disable=too-many-return-statements
def _raw_event_to_chunk(
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,
) -> 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``.
Returns ``None`` for events that should be silently skipped.
"""
event_type = raw.get("type")
# --- Message-level lifecycle ----------------------------------------
if event_type == "message_start":
message = raw.get("message", {})
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,
session_id=session_id,
chunk="",
output_tokens=usage.get("output_tokens"),
metadata={"stop_reason": delta.get("stop_reason")},
)
if event_type == "message_stop":
return cls._chunk(ChunkEnum.REPLY_END, session_id=session_id, chunk="")
# --- Content-block lifecycle ----------------------------------------
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_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.TOOL_CALL,
block_id=bid,
tool_call_id=content_block.get("id"),
tool_call_name=content_block.get("name"),
chunk=json.dumps(payload),
)
return None
if event_type == "content_block_delta":
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
if delta_type == "text_delta":
return cls._chunk(ChunkEnum.CONTENT, block_id=bid, chunk=delta.get("text", ""))
if delta_type == "thinking_delta":
return cls._chunk(ChunkEnum.THINK, block_id=bid, chunk=delta.get("thinking", ""))
if delta_type == "input_json_delta":
return cls._chunk(
ChunkEnum.TOOL_CALL,
block_id=bid,
tool_call_id=bid,
tool_call_name=tc_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
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":
return cls._chunk(ChunkEnum.THINK, block_id=bid, chunk="")
# text or unknown -> CONTENT
return cls._chunk(ChunkEnum.CONTENT, block_id=bid, chunk="")
# Ping / other unknown types -> skip
return None
@classmethod
def _message_content_to_chunks(
cls,
msg: Any,
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.
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):
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 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),
}
if hasattr(block, "is_error"):
payload["is_error"] = getattr(block, "is_error")
chunks.append(
cls._chunk(
ChunkEnum.TOOL_RESULT,
session_id=session_id,
block_id=tool_use_id,
tool_call_id=tool_use_id,
chunk=payload,
),
)
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
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)
# ----- reply / reply_stream --------------------------------------------
async def reply(self, inputs: Any, **kwargs) -> dict:
from claude_agent_sdk import query, ResultMessage
kwargs = self._merged_kwargs(kwargs)
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
if last_msg is None:
raise ValueError("No message received from Claude Code.")
result = {
"session_id": last_msg.session_id or "",
"last_message": asdict(last_msg),
"result": last_msg.result,
}
if kwargs.get("output_schema") is not None:
result["structured_output"] = last_msg.structured_output
return result
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
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] = {}
visible_tool_call_ids: set[str] = set()
current_session_id: str | None = None
emitted_content = False
received_result_message = False
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:
emitted_content = True
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):
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):
yield self._chunk(
ChunkEnum.ERROR,
session_id=msg.session_id or "",
chunk=str(msg.errors) if msg.errors else "Unknown error",
)
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:
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}")