diff --git a/llm_connect/claude_code.py b/llm_connect/claude_code.py index 0b0854e..ab1c3cf 100644 --- a/llm_connect/claude_code.py +++ b/llm_connect/claude_code.py @@ -2,12 +2,12 @@ Claude Code CLI adapter — runs the ``claude`` CLI as a subprocess. """ -import asyncio -import json -import os -import subprocess -from pathlib import Path -from typing import Optional +import asyncio +import json +import os +import subprocess +from pathlib import Path +from typing import Optional from llm_connect.adapter import LLMAdapter from llm_connect.models import RunConfig, LLMResponse @@ -26,23 +26,23 @@ class ClaudeCodeAdapter(LLMAdapter): length limits (compiled prompts can exceed 30 KB). """ - def __init__( - self, - cli_path: Optional[str] = None, - model: Optional[str] = None, - config: Optional[LLMConfig] = None, - ): - self._config = config or LLMConfig(provider="claude-code") - self._cli_path = cli_path or self._resolve_cli_path() - self._model = model + def __init__( + self, + cli_path: Optional[str] = None, + model: Optional[str] = None, + config: Optional[LLMConfig] = None, + ): + self._config = config or LLMConfig(provider="claude-code") + self._cli_path = cli_path or self._resolve_cli_path() + self._model = model # ── LLMAdapter interface ──────────────────────────────────────── - def execute_prompt(self, prompt: str, config: RunConfig) -> LLMResponse: - self._preflight_budget(config) - cmd = self._build_command(config) - - timeout = config.timeout_seconds or self._config.timeout_seconds + def execute_prompt(self, prompt: str, config: RunConfig) -> LLMResponse: + self._preflight_budget(config) + cmd = self._build_command(config) + + timeout = config.timeout_seconds or self._config.timeout_seconds try: result = subprocess.run( @@ -65,7 +65,7 @@ class ClaudeCodeAdapter(LLMAdapter): stderr=result.stderr, ) - content = result.stdout + content = _unwrap_cli_json_envelope(result.stdout, config) prompt_tokens = estimate_tokens(prompt) completion_tokens = estimate_tokens(content) @@ -86,10 +86,10 @@ class ClaudeCodeAdapter(LLMAdapter): self._consume_budget(config, response) return response - async def async_execute_prompt(self, prompt: str, config: RunConfig) -> LLMResponse: - """Native async implementation using asyncio.create_subprocess_exec.""" - self._preflight_budget(config) - cmd = self._build_command(config) + async def async_execute_prompt(self, prompt: str, config: RunConfig) -> LLMResponse: + """Native async implementation using asyncio.create_subprocess_exec.""" + self._preflight_budget(config) + cmd = self._build_command(config) timeout = config.timeout_seconds or self._config.timeout_seconds @@ -117,7 +117,7 @@ class ClaudeCodeAdapter(LLMAdapter): stderr=stderr_bytes.decode(), ) - content = stdout_bytes.decode() + content = _unwrap_cli_json_envelope(stdout_bytes.decode(), config) prompt_tokens = estimate_tokens(prompt) completion_tokens = estimate_tokens(content) @@ -139,49 +139,92 @@ class ClaudeCodeAdapter(LLMAdapter): self._consume_budget(config, response) return response - def validate_config(self, config: RunConfig) -> bool: - try: - result = subprocess.run( - [self._cli_path, "--version"], + def validate_config(self, config: RunConfig) -> bool: + try: + result = subprocess.run( + [self._cli_path, "--version"], capture_output=True, text=True, timeout=10, ) return result.returncode == 0 - except (subprocess.TimeoutExpired, FileNotFoundError, OSError): - return False - - def _build_command(self, config: RunConfig) -> list[str]: - cmd = [self._cli_path, "--print"] - if self._model: - cmd.extend(["--model", self._model]) - - json_schema = _json_schema_arg(config) - if json_schema: - cmd.extend(["--json-schema", json_schema]) - return cmd - - def _resolve_cli_path(self) -> str: - configured = ( - os.environ.get("LLM_CONNECT_CLAUDE_CLI_PATH") - or os.environ.get("CLAUDE_CLI_PATH") - or self._config.claude_cli_path - ) - if configured and configured != "claude": - return configured - - local_cli = Path.home() / ".local" / "bin" / "claude" - if local_cli.exists(): - return str(local_cli) - return configured or "claude" - - -def _json_schema_arg(config: RunConfig) -> str | None: - schema = (config.model_params or {}).get("json_schema") - if not schema: - return None - if isinstance(schema, str): - return schema - if isinstance(schema, dict): - return json.dumps(schema, separators=(",", ":")) - return None + except (subprocess.TimeoutExpired, FileNotFoundError, OSError): + return False + + def _build_command(self, config: RunConfig) -> list[str]: + cmd = [self._cli_path, "--print"] + if self._model: + cmd.extend(["--model", self._model]) + + json_schema = _json_schema_arg(config) + if json_schema: + cmd.extend(["--json-schema", json_schema]) + # With --json-schema alone the CLI prints conversational text on + # stdout while the structured payload ships on a sidecar channel + # callers cannot reach. --output-format json forces the structured + # response (wrapped in an envelope) onto stdout. + cmd.extend(["--output-format", "json"]) + return cmd + + def _resolve_cli_path(self) -> str: + configured = ( + os.environ.get("LLM_CONNECT_CLAUDE_CLI_PATH") + or os.environ.get("CLAUDE_CLI_PATH") + or self._config.claude_cli_path + ) + if configured and configured != "claude": + return configured + + local_cli = Path.home() / ".local" / "bin" / "claude" + if local_cli.exists(): + return str(local_cli) + return configured or "claude" + + +def _json_schema_arg(config: RunConfig) -> str | None: + schema = (config.model_params or {}).get("json_schema") + if not schema: + return None + if isinstance(schema, str): + return schema + if isinstance(schema, dict): + return json.dumps(schema, separators=(",", ":")) + return None + + +# Field names Claude Code's `--output-format json` envelope is known to use +# for the model's primary textual response. Probed in order; the first match +# wins. If none match (because the envelope shape is something we haven't +# seen), we return the raw envelope string so the caller still gets the data +# and can introspect it. +_ENVELOPE_TEXT_FIELDS = ("result", "result_text", "content", "text", "output") + + +def _unwrap_cli_json_envelope(stdout: str, config: RunConfig) -> str: + """Extract the model's payload from Claude CLI's --output-format json envelope. + + Only attempts unwrap when --json-schema was set, because that's the only + code path that adds --output-format json to the CLI invocation. Other + paths keep raw stdout (current behavior preserved). + """ + if not _json_schema_arg(config): + return stdout + text = stdout.strip() + if not text: + return stdout + try: + envelope = json.loads(text) + except json.JSONDecodeError: + return stdout + if not isinstance(envelope, dict): + return stdout + for key in _ENVELOPE_TEXT_FIELDS: + if key in envelope: + value = envelope[key] + if isinstance(value, str): + return value + if isinstance(value, (dict, list)): + return json.dumps(value) + # Unknown envelope shape — surface it raw so the operator can see it + # in the validation error and we can update _ENVELOPE_TEXT_FIELDS. + return stdout diff --git a/tests/test_claude_code.py b/tests/test_claude_code.py index d359a24..69dd309 100644 --- a/tests/test_claude_code.py +++ b/tests/test_claude_code.py @@ -39,15 +39,65 @@ def test_execute_prompt_passes_json_schema_to_claude_cli(monkeypatch): "--print", "--json-schema", '{"type":"object"}', + "--output-format", + "json", ], "input": "Produce a report.", "capture_output": True, "text": True, "timeout": 42, } + # Stdout shape that does not match any known envelope field is returned + # verbatim so the caller can introspect and we can extend the field list. assert response.content == '{"summary":"ok","recommendations":[]}' +def test_execute_prompt_unwraps_cli_json_envelope_result_field(monkeypatch): + """With --output-format json the CLI wraps the model payload in an + envelope. The adapter unwraps the textual result so the caller still + sees the model's structured-output JSON, not the envelope.""" + def fake_run(cmd, input, capture_output, text, timeout): # noqa: ANN001 + envelope = { + "type": "result", + "result": '{"summary":"ok","recommendations":[]}', + "total_cost_usd": 0.001, + } + import json as _json + return SimpleNamespace( + returncode=0, + stdout=_json.dumps(envelope), + stderr="", + ) + + monkeypatch.setattr("llm_connect.claude_code.subprocess.run", fake_run) + adapter = ClaudeCodeAdapter(cli_path="/custom/claude") + + response = adapter.execute_prompt( + "Produce a report.", + RunConfig(model_params={"json_schema": {"type": "object"}}), + ) + + assert response.content == '{"summary":"ok","recommendations":[]}' + + +def test_execute_prompt_no_unwrap_without_json_schema(monkeypatch): + """Without --json-schema we do not pass --output-format json, so the + envelope unwrap path stays inert and raw stdout passes through.""" + def fake_run(cmd, input, capture_output, text, timeout): # noqa: ANN001 + return SimpleNamespace( + returncode=0, + stdout='{"result":"this is just stdout, not an envelope"}', + stderr="", + ) + + monkeypatch.setattr("llm_connect.claude_code.subprocess.run", fake_run) + adapter = ClaudeCodeAdapter(cli_path="/custom/claude") + + response = adapter.execute_prompt("Plain prompt.", RunConfig()) + + assert response.content == '{"result":"this is just stdout, not an envelope"}' + + def test_claude_code_adapter_prefers_env_cli_path(monkeypatch): monkeypatch.setenv("LLM_CONNECT_CLAUDE_CLI_PATH", "/home/me/bin/claude")