mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-06 02:48:22 +00:00
* 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
657 lines
28 KiB
Python
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}")
|