458 lines
15 KiB
Python
458 lines
15 KiB
Python
# Claude Code CLI bridge
|
|
#
|
|
# Runs the locally installed `claude` binary in headless mode and hands back its
|
|
# answer, so a ComfyUI node can borrow the user's existing Claude Code login
|
|
# instead of a separate API key.
|
|
#
|
|
# The transport is Claude Code's bidirectional streaming JSON protocol:
|
|
#
|
|
# claude -p --input-format stream-json --output-format stream-json --verbose
|
|
#
|
|
# One JSON message goes in over stdin, a stream of JSON events comes back on
|
|
# stdout. That buys four things over passing a plain prompt:
|
|
#
|
|
# * Images travel as inline base64 content blocks. No scratch files, no Read
|
|
# tool, no permission grant, and one less turn per request.
|
|
# * Progress is visible while the model works, instead of a silent 45s block.
|
|
# * Rate-limit events are seen as they arrive, so quota problems are reported
|
|
# as quota problems.
|
|
# * The process can be killed the moment ComfyUI's queue is cancelled.
|
|
#
|
|
# Two constraints shape the rest. Windows caps a command line at 32,767
|
|
# characters and the H3 reference guide alone is ~40 KB, so a long system prompt
|
|
# is folded into the message rather than passed as an argument. And
|
|
# `--setting-sources ""` keeps the user's own CLAUDE.md, skills and hooks out of
|
|
# the request - a personal "always use Go best practices" rule has no business
|
|
# rewriting a video prompt.
|
|
|
|
import base64
|
|
import io
|
|
import json
|
|
import os
|
|
import queue
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
|
|
# Present only when running inside ComfyUI; the module is importable standalone.
|
|
try:
|
|
import comfy.model_management as _model_management
|
|
except Exception:
|
|
_model_management = None
|
|
|
|
# Model aliases the CLI accepts. A full id (claude-sonnet-5) works too.
|
|
CLAUDE_CODE_MODELS = ["sonnet", "opus", "haiku", "fable", "default"]
|
|
|
|
# Read-only tools, granted only when research mode is on. Bash/Write/Edit are
|
|
# deliberately absent - a prompt writer has no business changing anything.
|
|
RESEARCH_TOOLS = ["Read", "Glob", "Grep", "WebSearch", "WebFetch"]
|
|
|
|
# A system prompt longer than this moves into the stdin message instead of being
|
|
# passed as an argument, well under the Windows command-line limit.
|
|
_ARGV_SAFE_LIMIT = 8000
|
|
|
|
_EMBEDDED_SYSTEM_FRAMING = (
|
|
"You follow the specification in the user message exactly. Emit only what it "
|
|
"asks for: no preamble, no commentary, no summary of what you did."
|
|
)
|
|
|
|
# How often the reader loop wakes up to check for cancellation and timeout.
|
|
_TICK_SECONDS = 0.2
|
|
|
|
# Grace period between asking the CLI to stop and killing it outright.
|
|
_KILL_GRACE_SECONDS = 3
|
|
|
|
_cli_cache = {"path": None, "checked": False}
|
|
|
|
|
|
class ClaudeCodeError(RuntimeError):
|
|
"""The CLI is missing, timed out, or reported a failure."""
|
|
|
|
|
|
def find_cli(refresh=False):
|
|
"""Absolute path to the claude binary, or None when it is not installed."""
|
|
if _cli_cache["checked"] and not refresh:
|
|
return _cli_cache["path"]
|
|
|
|
candidates = []
|
|
override = os.environ.get("CLAUDE_CODE_PATH")
|
|
if override:
|
|
candidates.append(override)
|
|
|
|
found = shutil.which("claude")
|
|
if found:
|
|
candidates.append(found)
|
|
|
|
home = os.path.expanduser("~")
|
|
candidates.extend(
|
|
[
|
|
os.path.join(home, ".local", "bin", "claude.exe"),
|
|
os.path.join(home, ".local", "bin", "claude"),
|
|
os.path.join(os.environ.get("APPDATA", ""), "npm", "claude.cmd"),
|
|
"/usr/local/bin/claude",
|
|
"/opt/homebrew/bin/claude",
|
|
]
|
|
)
|
|
|
|
path = next((c for c in candidates if c and os.path.isfile(c)), None)
|
|
_cli_cache["path"] = path
|
|
_cli_cache["checked"] = True
|
|
return path
|
|
|
|
|
|
def is_available():
|
|
return find_cli() is not None
|
|
|
|
|
|
def is_interrupt(exc):
|
|
"""Did this exception come from the user cancelling the ComfyUI queue?"""
|
|
if _model_management is None:
|
|
return False
|
|
return isinstance(exc, _model_management.InterruptProcessingException)
|
|
|
|
|
|
def _processing_interrupted():
|
|
if _model_management is None:
|
|
return False
|
|
try:
|
|
return _model_management.processing_interrupted()
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def _subprocess_flags():
|
|
"""Keep Windows from flashing a console window on every node run."""
|
|
if sys.platform == "win32":
|
|
return {"creationflags": getattr(subprocess, "CREATE_NO_WINDOW", 0)}
|
|
return {}
|
|
|
|
|
|
def _child_env(use_subscription):
|
|
"""
|
|
Environment for the child process.
|
|
|
|
ANTHROPIC_API_KEY is dropped by default: with it set, Claude Code bills the
|
|
API account, which defeats the point of routing through a subscription seat.
|
|
The other nodes in this pack want that key, hence the per-call switch.
|
|
"""
|
|
env = os.environ.copy()
|
|
if use_subscription:
|
|
for name in ("ANTHROPIC_API_KEY", "CLAUDE_API_KEY", "ANTHROPIC_AUTH_TOKEN"):
|
|
env.pop(name, None)
|
|
return env
|
|
|
|
|
|
def _image_block(image):
|
|
"""PIL image -> an Anthropic base64 image content block."""
|
|
buffer = io.BytesIO()
|
|
image.convert("RGB").save(buffer, format="PNG")
|
|
encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
|
|
return {
|
|
"type": "image",
|
|
"source": {"type": "base64", "media_type": "image/png", "data": encoded},
|
|
}
|
|
|
|
|
|
def _user_message(text, images):
|
|
"""
|
|
The single stream-json message sent to the CLI.
|
|
|
|
Images come first: the model reads them before the instruction that refers
|
|
to them, which is what Anthropic recommends for image-led prompts.
|
|
"""
|
|
content = [_image_block(image) for image in images]
|
|
content.append({"type": "text", "text": text})
|
|
return {"type": "user", "message": {"role": "user", "content": content}}
|
|
|
|
|
|
def _pump(stream, sink, tag):
|
|
"""Drain a pipe on its own thread so the child never blocks on a full buffer."""
|
|
try:
|
|
for line in stream:
|
|
sink.put((tag, line))
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
sink.put((tag, None))
|
|
try:
|
|
stream.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _write_stdin(process, payload):
|
|
"""
|
|
Feed the request in on a thread.
|
|
|
|
A prompt carrying a few base64 images is far larger than the OS pipe buffer,
|
|
so writing it inline would deadlock against a child that has not started
|
|
reading yet.
|
|
"""
|
|
try:
|
|
process.stdin.write(payload)
|
|
process.stdin.flush()
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
try:
|
|
process.stdin.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _stop(process):
|
|
"""Ask the CLI to exit, then insist."""
|
|
if process.poll() is not None:
|
|
return
|
|
try:
|
|
process.terminate()
|
|
process.wait(timeout=_KILL_GRACE_SECONDS)
|
|
except Exception:
|
|
try:
|
|
process.kill()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _describe_event(event):
|
|
"""A short progress line for an event, or None when it is not worth showing."""
|
|
kind = event.get("type")
|
|
|
|
if kind == "system":
|
|
# thinking_tokens arrives many times a second; init is the useful one.
|
|
if event.get("subtype") == "init":
|
|
return f"session {(event.get('session_id') or '')[:8]} ready"
|
|
return None
|
|
|
|
if kind == "assistant":
|
|
parts = []
|
|
for block in event.get("message", {}).get("content", []):
|
|
if block.get("type") == "tool_use":
|
|
parts.append(f"using {block.get('name', 'a tool')}")
|
|
elif block.get("type") == "text":
|
|
snippet = " ".join((block.get("text") or "").split())
|
|
if snippet:
|
|
parts.append(snippet[:110] + ("..." if len(snippet) > 110 else ""))
|
|
return " | ".join(parts) or None
|
|
|
|
return None
|
|
|
|
|
|
def _rate_limit_note(info):
|
|
"""A warning for a rate-limit event, or None while everything is fine."""
|
|
status = (info.get("status") or "").lower()
|
|
if status in ("", "allowed", "ok"):
|
|
return None
|
|
|
|
window = info.get("rateLimitType") or "usage"
|
|
note = f"Claude Code {window} limit: {status}"
|
|
resets_at = info.get("resetsAt")
|
|
if resets_at:
|
|
try:
|
|
when = time.strftime("%H:%M", time.localtime(int(resets_at)))
|
|
note += f" (resets around {when})"
|
|
except Exception:
|
|
pass
|
|
return note
|
|
|
|
|
|
def run_claude_code(
|
|
prompt,
|
|
system_prompt=None,
|
|
images=None,
|
|
model="sonnet",
|
|
timeout=600,
|
|
tools=None,
|
|
working_dir=None,
|
|
resume_session_id=None,
|
|
use_subscription=True,
|
|
on_progress=None,
|
|
add_dirs=None,
|
|
):
|
|
"""
|
|
Run one headless Claude Code turn and return its result plus telemetry.
|
|
|
|
`images` is a list of PIL images, sent inline. `tools` is an explicit
|
|
allow-list; None grants nothing. `add_dirs` are extra directories the
|
|
granted tools may read (`--add-dir`), on top of the working directory.
|
|
`on_progress` is called with short status strings while the turn runs.
|
|
Raises ClaudeCodeError on failure, and lets a ComfyUI cancellation
|
|
propagate as an interrupt.
|
|
"""
|
|
cli = find_cli()
|
|
if not cli:
|
|
raise ClaudeCodeError(
|
|
"The Claude Code CLI was not found. Install it from "
|
|
"https://claude.com/claude-code, or set CLAUDE_CODE_PATH to the binary."
|
|
)
|
|
|
|
images = list(images or [])
|
|
body = prompt
|
|
|
|
args = [
|
|
cli,
|
|
"-p",
|
|
"--input-format", "stream-json",
|
|
"--output-format", "stream-json",
|
|
"--verbose",
|
|
"--setting-sources", "",
|
|
]
|
|
|
|
if model and model != "default":
|
|
args += ["--model", model]
|
|
|
|
# Long system prompts cannot be arguments, so they ride along in the message.
|
|
if system_prompt and len(system_prompt) > _ARGV_SAFE_LIMIT:
|
|
args += ["--system-prompt", _EMBEDDED_SYSTEM_FRAMING]
|
|
body = (
|
|
"=== BEGIN SPECIFICATION - follow it exactly ===\n"
|
|
f"{system_prompt}\n"
|
|
"=== END SPECIFICATION ===\n\n"
|
|
f"{prompt}"
|
|
)
|
|
elif system_prompt:
|
|
args += ["--system-prompt", system_prompt]
|
|
|
|
if tools:
|
|
args += ["--allowedTools", ",".join(tools)]
|
|
# A headless run cannot answer a permission prompt; without this the
|
|
# tool call is denied and the model answers blind.
|
|
args += ["--permission-mode", "dontAsk"]
|
|
|
|
for extra in add_dirs or []:
|
|
if extra and os.path.isdir(extra):
|
|
args += ["--add-dir", extra]
|
|
|
|
if resume_session_id:
|
|
args += ["--resume", resume_session_id]
|
|
|
|
# An isolated working directory keeps the CLI out of the ComfyUI tree.
|
|
scratch = tempfile.mkdtemp(prefix="apnext_claude_code_")
|
|
cwd = working_dir if working_dir and os.path.isdir(working_dir) else scratch
|
|
|
|
payload = json.dumps(_user_message(body, images), ensure_ascii=False) + "\n"
|
|
|
|
events = queue.Queue()
|
|
result_payload = None
|
|
rate_limit_info = {}
|
|
stderr_lines = []
|
|
warned_rate_limit = False
|
|
open_streams = 2
|
|
|
|
try:
|
|
try:
|
|
process = subprocess.Popen(
|
|
args,
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
encoding="utf-8",
|
|
errors="replace",
|
|
bufsize=1,
|
|
cwd=cwd,
|
|
env=_child_env(use_subscription),
|
|
**_subprocess_flags(),
|
|
)
|
|
except OSError as exc:
|
|
raise ClaudeCodeError(f"Could not start the Claude Code CLI at {cli}: {exc}")
|
|
|
|
threading.Thread(target=_write_stdin, args=(process, payload), daemon=True).start()
|
|
threading.Thread(target=_pump, args=(process.stdout, events, "out"), daemon=True).start()
|
|
threading.Thread(target=_pump, args=(process.stderr, events, "err"), daemon=True).start()
|
|
|
|
deadline = time.monotonic() + timeout
|
|
|
|
while open_streams:
|
|
try:
|
|
tag, line = events.get(timeout=_TICK_SECONDS)
|
|
except queue.Empty:
|
|
if _processing_interrupted():
|
|
_stop(process)
|
|
# Raises ComfyUI's own interrupt so the queue stops cleanly.
|
|
_model_management.throw_exception_if_processing_interrupted()
|
|
if time.monotonic() > deadline:
|
|
_stop(process)
|
|
raise ClaudeCodeError(
|
|
f"Claude Code did not finish within {timeout}s. Raise the timeout, "
|
|
"or use a faster model."
|
|
)
|
|
continue
|
|
|
|
if line is None:
|
|
open_streams -= 1
|
|
continue
|
|
|
|
if tag == "err":
|
|
text = line.strip()
|
|
if text:
|
|
stderr_lines.append(text)
|
|
continue
|
|
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
|
|
try:
|
|
event = json.loads(line)
|
|
except ValueError:
|
|
continue
|
|
|
|
kind = event.get("type")
|
|
|
|
if kind == "result":
|
|
result_payload = event
|
|
continue
|
|
|
|
if kind == "rate_limit_event":
|
|
rate_limit_info = event.get("rate_limit_info") or {}
|
|
note = _rate_limit_note(rate_limit_info)
|
|
if note and not warned_rate_limit:
|
|
warned_rate_limit = True
|
|
print(f"⚠️ {note}")
|
|
continue
|
|
|
|
if on_progress:
|
|
description = _describe_event(event)
|
|
if description:
|
|
on_progress(description)
|
|
|
|
process.wait(timeout=_KILL_GRACE_SECONDS)
|
|
|
|
if result_payload is None:
|
|
detail = " ".join(stderr_lines) or f"exit code {process.returncode}"
|
|
raise ClaudeCodeError(f"Claude Code returned no result: {detail}")
|
|
|
|
text = (result_payload.get("result") or "").strip()
|
|
|
|
if result_payload.get("is_error") or result_payload.get("subtype") not in (None, "success"):
|
|
detail = text or " ".join(stderr_lines) or result_payload.get("subtype") or "unknown error"
|
|
note = _rate_limit_note(rate_limit_info)
|
|
if note:
|
|
detail = f"{detail} ({note})"
|
|
if "login" in detail.lower() or "authenticat" in detail.lower():
|
|
detail += (
|
|
" - run `claude` once in a terminal to log in, or untick the "
|
|
"subscription switch to use ANTHROPIC_API_KEY."
|
|
)
|
|
raise ClaudeCodeError(f"Claude Code reported an error: {detail}")
|
|
|
|
if not text:
|
|
raise ClaudeCodeError("Claude Code returned an empty result.")
|
|
|
|
return {
|
|
"text": text,
|
|
"session_id": result_payload.get("session_id") or "",
|
|
"cost_usd": float(result_payload.get("total_cost_usd") or 0.0),
|
|
"duration_ms": int(result_payload.get("duration_ms") or 0),
|
|
"num_turns": int(result_payload.get("num_turns") or 0),
|
|
"model": model,
|
|
"rate_limit": rate_limit_info,
|
|
}
|
|
|
|
finally:
|
|
shutil.rmtree(scratch, ignore_errors=True)
|