From b9d0aec830d41ada4e9faa6fd415824d4b09e2f1 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Thu, 2 Jul 2026 18:47:44 -0700 Subject: [PATCH 1/2] feat: arm-B CLI micro-edit proxy bench runner over slots/set-slot/vary + CQL (BE-2309) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `comfy_cli/bench/run_arm_b.py`: a minimal Claude Opus (claude-opus-4-8, Anthropic Messages API) agent loop that exposes the existing `comfy workflow` micro-edit substrate as tools — read=slots+cql, edit=set-slot, produce-variants=vary — and drives the shared BE-2302 t1–t4 task matrix (vendored byte-identical from child 1's tasks.mjs into bench/tasks.json). The loop reads/edits a temp frontend-format workflow JSON on disk (no full-JSON round-trip to the model), captures the Anthropic usage block per API call plus tool-call count and per-call payload bytes, and prices each turn with the same table as comfy-inapp-agent/agent-server/usage.mjs so the two arms are comparable. It emits per-(task, turn) NDJSON shaped identically to arm A's arm-a.ndjson (verified: child 1's report.mjs consumes bench/out/arm-b.ndjson directly). Spike-sanctioned PROXY for Kishore's not-yet-accessible CLI-agent prototype: every row is stamped `"proxy": true`, and README + module docstring document the narrow swap-in seam (inject the client / replace Driver, keep ToolDispatcher + build_row). Offline-first: committed object_info + seed fixtures run t1/t2/t4 without a live ComfyUI; t3 (~150-node build) is recorded as a RESULT ceiling (CQL/object_info coverage limit). Unit tests stub the model client (no live calls, CI stays offline); a `bench` extra + README cover the live run. --- comfy_cli/bench/README.md | 101 ++ comfy_cli/bench/__init__.py | 1 + .../bench/fixtures/edit_session_seed.json | 546 +++++++++++ comfy_cli/bench/fixtures/object_info.json | 467 +++++++++ comfy_cli/bench/fixtures/txt2img_seed.json | 547 +++++++++++ comfy_cli/bench/out/.gitignore | 4 + comfy_cli/bench/out/.gitkeep | 1 + comfy_cli/bench/run_arm_b.py | 926 ++++++++++++++++++ comfy_cli/bench/tasks.json | 39 + pyproject.toml | 5 + tests/comfy_cli/bench/__init__.py | 0 tests/comfy_cli/bench/test_run_arm_b.py | 282 ++++++ 12 files changed, 2919 insertions(+) create mode 100644 comfy_cli/bench/README.md create mode 100644 comfy_cli/bench/__init__.py create mode 100644 comfy_cli/bench/fixtures/edit_session_seed.json create mode 100644 comfy_cli/bench/fixtures/object_info.json create mode 100644 comfy_cli/bench/fixtures/txt2img_seed.json create mode 100644 comfy_cli/bench/out/.gitignore create mode 100644 comfy_cli/bench/out/.gitkeep create mode 100644 comfy_cli/bench/run_arm_b.py create mode 100644 comfy_cli/bench/tasks.json create mode 100644 tests/comfy_cli/bench/__init__.py create mode 100644 tests/comfy_cli/bench/test_run_arm_b.py diff --git a/comfy_cli/bench/README.md b/comfy_cli/bench/README.md new file mode 100644 index 00000000..f8d565e8 --- /dev/null +++ b/comfy_cli/bench/README.md @@ -0,0 +1,101 @@ +# Arm-B bench runner — CLI micro-edit **proxy** (BE-2309 / BE-2302) + +> ⚠️ **This is a spike-sanctioned PROXY, not a product feature and not the real thing.** +> Arm B of the BE-2302 A/B benchmark is meant to measure *Kishore's CLI-agent prototype*, +> which is not yet accessible. Until it is, this runner stands in for it by driving Claude +> over the `comfy workflow` micro-edit substrate that already lives in this repo. Every +> emitted NDJSON row is stamped `"proxy": true` so no downstream reader mistakes proxy +> numbers for the real prototype's. See **Swapping in the real prototype** below. + +## What it is + +`run_arm_b.py` is a minimal agent loop (Claude Opus `claude-opus-4-8`, raw Anthropic +Messages API) that exposes exactly the arm-B toolset BE-2302 defines — "4–5 generic +micro-edit tools operating on a temp JSON workflow on disk": + +| Arm-B role | Tool | comfy-cli command / layer it wraps | +| ---------------- | ---------- | ----------------------------------------------------------- | +| read | `slots` | `comfy workflow slots` (`command/workflow.py`) | +| read (validate) | `cql` | the `comfy_cli/cql/` node/type catalog (`Graph`, `engine.py`) | +| edit | `set_slot` | `comfy workflow set-slot` | +| produce-variants | `vary` | `comfy workflow vary` | + +The loop reads/edits a **temp frontend-format workflow JSON on disk** — the model only ever +sees compact slot manifests, never the full workflow JSON (no full-JSON round-trip). It runs +the same **t1–t4** task prompts as arm A, vendored byte-identical from child 1's +`comfy-inapp-agent/agent-server/bench/tasks.mjs` into [`tasks.json`](./tasks.json). + +## Telemetry & output + +Per Anthropic API call it captures the `usage` block (`input_tokens`, `output_tokens`, +`cache_read_input_tokens`, `cache_creation_input_tokens`) plus tool-call count and per-call +payload bytes, and folds a task's turns into one NDJSON row per `(task, turn)`. Cost is +recomputed from the token counts with the **same pricing table** as +`comfy-inapp-agent/agent-server/usage.mjs` (`PRICING`): Opus 4.8 = $5/M in, $25/M out, cache +read 0.1×, cache write 1.25× — so arm A and arm B are dollar-comparable. + +Rows are written to [`out/arm-b.ndjson`](./out/), **shaped identically to arm A's +`arm-a.ndjson`**, so child 1's `report.mjs` consumes them directly: + +```bash +# in comfy-inapp-agent/agent-server/ +node bench/report.mjs --arm-b /path/to/comfy-cli/comfy_cli/bench/out/arm-b.ndjson +``` + +## Offline dry run (no API key, no network) — the acceptance path + +```bash +python -m comfy_cli.bench.run_arm_b --dry-run +``` + +Drives a **stubbed, deterministic** model against the committed +[`fixtures/object_info.json`](./fixtures/) and seed workflows, exercising the full tool +dispatch + usage-parsing + NDJSON pipeline, and writes a well-formed `out/arm-b.ndjson`. This +is what CI and the unit tests (`tests/comfy_cli/bench/test_run_arm_b.py`) run — **no live +calls**. t1/t2/t4 run fully offline against the fixture. + +**t3 is the documented ceiling.** Arm B's slots/set-slot/vary toolset can only *edit* an +existing graph, and the committed `object_info` covers only the 6 base SD1.5 node classes +(no LoRA/ControlNet/refiner/upscaler/face-detailer). A ~150-node *build* is therefore out of +reach; the runner records this as a **RESULT** — the t3 row carries `"outcome": "ceiling"` +and a `note` explaining the CQL/object_info coverage limit — rather than pretending to +succeed. + +## Live run + +```bash +pip install -e '.[bench]' # adds the `anthropic` SDK +export ANTHROPIC_API_KEY=sk-ant-... +python -m comfy_cli.bench.run_arm_b --live +# optional: --only t2 t4 --model claude-opus-4-8 --out out/arm-b.ndjson +# --object-info (point at a live object_info dump for wider node coverage) +``` + +A live run replaces the stub with a real `anthropic.Anthropic` client; everything else — tool +dispatch, telemetry, NDJSON shape — is identical to the dry run. For wider node coverage than +the committed SD1.5 fixture (e.g. to attempt t3 for real), pass `--object-info` a dump saved +from a live server: `comfy --json ... > object_info.json`, or point `slots`/`set-slot` at a +running ComfyUI. + +## Swapping in the real prototype (the documented seam) + +The proxy is built so Kishore's real CLI-agent prototype drops into one of two narrow seams, +with the telemetry + NDJSON shape unchanged: + +1. **Replace the model client.** The client is injected into `Driver` and only needs a + `messages.create(...)` returning `.content` / `.stop_reason` / `.usage`. `build_live_client()` + returns the real Anthropic SDK; point it at the prototype's endpoint instead. +2. **Replace the driver.** For a prototype that speaks its own protocol, implement a driver in + place of `Driver.run_turn` that keeps calling `ToolDispatcher` (the comfy-cli substrate) and + emits rows via `build_row(...)`. The pricing, usage-parsing, ceiling handling, and NDJSON + contract all stay put. + +When the real prototype lands, drop the `"proxy": true` stamp in `build_row` (and update this +README) so the rows advertise themselves as the genuine article. + +## Keeping the task matrix in sync + +`tasks.json` is a **vendored copy** of child 1's `tasks.mjs` prompt strings and MUST stay +byte-identical for the two arms to be comparable. It was generated by evaluating `tasks.mjs` +and dumping its `TASKS` prompts, so re-vendor it (don't hand-edit) if the shared matrix +changes upstream. diff --git a/comfy_cli/bench/__init__.py b/comfy_cli/bench/__init__.py new file mode 100644 index 00000000..da9c62fe --- /dev/null +++ b/comfy_cli/bench/__init__.py @@ -0,0 +1 @@ +"""BE-2302 A/B benchmark — arm B (CLI micro-edit proxy) runner. See run_arm_b.py + README.md.""" diff --git a/comfy_cli/bench/fixtures/edit_session_seed.json b/comfy_cli/bench/fixtures/edit_session_seed.json new file mode 100644 index 00000000..603184fc --- /dev/null +++ b/comfy_cli/bench/fixtures/edit_session_seed.json @@ -0,0 +1,546 @@ +{ + "id": "2ba0b800-2f13-4f21-b8d6-c6cdb0152cae", + "revision": 0, + "last_node_id": 16, + "last_link_id": 9, + "nodes": [ + { + "id": 4, + "type": "CheckpointLoaderSimple", + "pos": [ + 10, + 300 + ], + "size": [ + 320, + 154.65625 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "slot_index": 0, + "links": [ + 1 + ] + }, + { + "name": "CLIP", + "type": "CLIP", + "slot_index": 1, + "links": [ + 3, + 5 + ] + }, + { + "name": "VAE", + "type": "VAE", + "slot_index": 2, + "links": [ + 8 + ] + } + ], + "properties": { + "Node name for S&R": "CheckpointLoaderSimple", + "cnr_id": "comfy-core", + "ver": "0.3.65", + "models": [ + { + "name": "v1-5-pruned-emaonly-fp16.safetensors", + "url": "https://huggingface.co/Comfy-Org/stable-diffusion-v1-5-archive/resolve/main/v1-5-pruned-emaonly-fp16.safetensors?download=true", + "directory": "checkpoints" + } + ] + }, + "widgets_values": [ + "v1-5-pruned-emaonly-fp16.safetensors" + ] + }, + { + "id": 3, + "type": "KSampler", + "pos": [ + 920, + 170 + ], + "size": [ + 320, + 480 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 1 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 4 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 6 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 2 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "slot_index": 0, + "links": [ + 7 + ] + } + ], + "properties": { + "Node name for S&R": "KSampler", + "cnr_id": "comfy-core", + "ver": "0.3.65" + }, + "widgets_values": [ + 123456789, + 20, + 7.0, + "euler", + "normal", + 1.0 + ] + }, + { + "id": 8, + "type": "VAEDecode", + "pos": [ + 1000, + 710 + ], + "size": [ + 225, + 96 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": 8 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "slot_index": 0, + "links": [ + 9 + ] + } + ], + "properties": { + "Node name for S&R": "VAEDecode", + "cnr_id": "comfy-core", + "ver": "0.3.65" + }, + "widgets_values": [] + }, + { + "id": 9, + "type": "SaveImage", + "pos": [ + 1270, + 170 + ], + "size": [ + 470, + 560 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 9 + } + ], + "outputs": [], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.65" + }, + "widgets_values": [ + "landscape" + ] + }, + { + "id": 7, + "type": "CLIPTextEncode", + "pos": [ + 430, + 530 + ], + "size": [ + 420, + 170 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 5 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 6 + ] + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode", + "cnr_id": "comfy-core", + "ver": "0.3.65" + }, + "widgets_values": [ + "text, watermark" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 5, + "type": "EmptyLatentImage", + "pos": [ + 490, + 900 + ], + "size": [ + 320, + 168 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "slot_index": 0, + "links": [ + 2 + ] + } + ], + "properties": { + "Node name for S&R": "EmptyLatentImage", + "cnr_id": "comfy-core", + "ver": "0.3.65" + }, + "widgets_values": [ + 512, + 512, + 1 + ] + }, + { + "id": 6, + "type": "CLIPTextEncode", + "pos": [ + 420, + 220 + ], + "size": [ + 430, + 260 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 3 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 4 + ] + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode", + "cnr_id": "comfy-core", + "ver": "0.3.65" + }, + "widgets_values": [ + "a serene mountain landscape at golden hour, misty valleys" + ], + "color": "#232", + "bgcolor": "#353" + }, + { + "id": 15, + "type": "MarkdownNote", + "pos": [ + 400, + -320 + ], + "size": [ + 470, + 430 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "Note: Prompt", + "properties": {}, + "widgets_values": [ + "**Prompts usually have two types:**\n\n* **Positive Prompt:** Tells the model what you *want* to see.\n* **Negative Prompt:** Tells the model what you *don\u2019t want* to see.\n\nThink of it as:
\n\ud83d\udc49 Positive = \u201cDo this\u201d
\n\ud83d\udc49 Negative = \u201cDon\u2019t do this\u201d\n\n\nDifferent models may interpret prompts differently.
\nSome prefer short, simple phrases; others respond well to detailed descriptions or styles.\nExperiment to see how each model reacts.\n\nAbout SD1.5:
\nStable Diffusion 1.5 is one of the most popular base models.\nIt works best with short, clear prompts and simple concepts, and it has a natural, realistic visual style.\n" + ], + "color": "#432", + "bgcolor": "#000" + }, + { + "id": 14, + "type": "MarkdownNote", + "pos": [ + 1270, + 780 + ], + "size": [ + 470, + 130 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "Note: Output", + "properties": {}, + "widgets_values": [ + "Image will auto-save to the `ComfyUI/output` folder. You can also right-click on the `Save Image` node and then use the menu to save the image locally." + ], + "color": "#432", + "bgcolor": "#000" + }, + { + "id": 13, + "type": "MarkdownNote", + "pos": [ + 460, + 1180 + ], + "size": [ + 330, + 163.953125 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "Note: Image size", + "properties": {}, + "widgets_values": [ + "Different models are trained based on datasets of different image sizes. This workflow is using Stable Diffusion 1.5, which is trained based on 512x512 datasets. So, it doesn't perform well on large image sizes." + ], + "color": "#432", + "bgcolor": "#000" + }, + { + "id": 11, + "type": "MarkdownNote", + "pos": [ + -470, + 160 + ], + "size": [ + 400, + 530.890625 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "Note: Model link", + "properties": {}, + "widgets_values": [ + "[Tutorial](https://docs.comfy.org/tutorials/basic/text-to-image)\n\n\n## Model link\n\n**checkpoints**\n\n- [v1-5-pruned-emaonly-fp16.safetensors](https://huggingface.co/Comfy-Org/stable-diffusion-v1-5-archive/resolve/main/v1-5-pruned-emaonly-fp16.safetensors?download=true)\n\n\nModel Storage Location\n\n```\n\ud83d\udcc2 ComfyUI/\n\u251c\u2500\u2500 \ud83d\udcc2 models/\n\u2502 \u2514\u2500\u2500 \ud83d\udcc2 checkpoints/\n\u2502 \u2514\u2500\u2500 v1-5-pruned-emaonly-fp16.safetensors\n```\n\n\n## Report issue\nIf you have any problems running this workflow, please report template-related issues via this link: [report the template issue here](https://github.com/Comfy-Org/workflow_templates/issues)" + ], + "color": "#432", + "bgcolor": "#000" + } + ], + "links": [ + [ + 1, + 4, + 0, + 3, + 0, + "MODEL" + ], + [ + 2, + 5, + 0, + 3, + 3, + "LATENT" + ], + [ + 3, + 4, + 1, + 6, + 0, + "CLIP" + ], + [ + 4, + 6, + 0, + 3, + 1, + "CONDITIONING" + ], + [ + 5, + 4, + 1, + 7, + 0, + "CLIP" + ], + [ + 6, + 7, + 0, + 3, + 2, + "CONDITIONING" + ], + [ + 7, + 3, + 0, + 8, + 0, + "LATENT" + ], + [ + 8, + 4, + 2, + 8, + 1, + "VAE" + ], + [ + 9, + 8, + 0, + 9, + 0, + "IMAGE" + ] + ], + "groups": [ + { + "id": 1, + "title": "Step 1 - Load model", + "bounding": [ + -40, + 130, + 420, + 470 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + }, + { + "id": 2, + "title": "Step 3 - Image size", + "bounding": [ + 400, + 800, + 480, + 310 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + }, + { + "id": 3, + "title": "Step 2 - Prompt", + "bounding": [ + 400, + 130, + 480, + 640 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + } + ], + "config": {}, + "extra": { + "ds": { + "scale": 0.5131581182307069, + "offset": [ + 979.5226642853634, + 273.924658465434 + ] + }, + "frontendVersion": "1.42.15", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/comfy_cli/bench/fixtures/object_info.json b/comfy_cli/bench/fixtures/object_info.json new file mode 100644 index 00000000..af497615 --- /dev/null +++ b/comfy_cli/bench/fixtures/object_info.json @@ -0,0 +1,467 @@ +{ + "CLIPTextEncode": { + "input": { + "required": { + "text": [ + "STRING", + { + "multiline": true, + "dynamicPrompts": true, + "tooltip": "The text to be encoded." + } + ], + "clip": [ + "CLIP", + { + "tooltip": "The CLIP model used for encoding the text." + } + ] + } + }, + "input_order": { + "required": [ + "text", + "clip" + ] + }, + "is_input_list": false, + "output": [ + "CONDITIONING" + ], + "output_is_list": [ + false + ], + "output_name": [ + "CONDITIONING" + ], + "name": "CLIPTextEncode", + "display_name": "CLIP Text Encode (Prompt)", + "description": "Encodes a text prompt using a CLIP model into an embedding that can be used to guide the diffusion model towards generating specific images.", + "python_module": "nodes", + "category": "conditioning", + "output_node": false, + "has_intermediate_output": false, + "output_tooltips": [ + "A conditioning containing the embedded text used to guide the diffusion model." + ], + "search_aliases": [ + "text", + "prompt", + "text prompt", + "positive prompt", + "negative prompt", + "encode text", + "text encoder", + "encode prompt" + ] + }, + "CheckpointLoaderSimple": { + "input": { + "required": { + "ckpt_name": [ + [ + "sd_xl_turbo_1.0_fp16.safetensors", + "v1-5-pruned-emaonly-fp16.safetensors" + ], + { + "tooltip": "The name of the checkpoint (model) to load." + } + ] + } + }, + "input_order": { + "required": [ + "ckpt_name" + ] + }, + "is_input_list": false, + "output": [ + "MODEL", + "CLIP", + "VAE" + ], + "output_is_list": [ + false, + false, + false + ], + "output_name": [ + "MODEL", + "CLIP", + "VAE" + ], + "name": "CheckpointLoaderSimple", + "display_name": "Load Checkpoint", + "description": "Loads a diffusion model checkpoint, diffusion models are used to denoise latents.", + "python_module": "nodes", + "category": "loaders", + "output_node": false, + "has_intermediate_output": false, + "output_tooltips": [ + "The model used for denoising latents.", + "The CLIP model used for encoding text prompts.", + "The VAE model used for encoding and decoding images to and from latent space." + ], + "search_aliases": [ + "load model", + "checkpoint", + "model loader", + "load checkpoint", + "ckpt", + "model" + ] + }, + "EmptyLatentImage": { + "input": { + "required": { + "width": [ + "INT", + { + "default": 512, + "min": 16, + "max": 16384, + "step": 8, + "tooltip": "The width of the latent images in pixels." + } + ], + "height": [ + "INT", + { + "default": 512, + "min": 16, + "max": 16384, + "step": 8, + "tooltip": "The height of the latent images in pixels." + } + ], + "batch_size": [ + "INT", + { + "default": 1, + "min": 1, + "max": 4096, + "tooltip": "The number of latent images in the batch." + } + ] + } + }, + "input_order": { + "required": [ + "width", + "height", + "batch_size" + ] + }, + "is_input_list": false, + "output": [ + "LATENT" + ], + "output_is_list": [ + false + ], + "output_name": [ + "LATENT" + ], + "name": "EmptyLatentImage", + "display_name": "Empty Latent Image", + "description": "Create a new batch of empty latent images to be denoised via sampling.", + "python_module": "nodes", + "category": "latent", + "output_node": false, + "has_intermediate_output": false, + "output_tooltips": [ + "The empty latent image batch." + ], + "search_aliases": [ + "empty", + "empty latent", + "new latent", + "create latent", + "blank latent", + "blank" + ] + }, + "KSampler": { + "input": { + "required": { + "model": [ + "MODEL", + { + "tooltip": "The model used for denoising the input latent." + } + ], + "seed": [ + "INT", + { + "default": 0, + "min": 0, + "max": 18446744073709551615, + "control_after_generate": true, + "tooltip": "The random seed used for creating the noise." + } + ], + "steps": [ + "INT", + { + "default": 20, + "min": 1, + "max": 10000, + "tooltip": "The number of steps used in the denoising process." + } + ], + "cfg": [ + "FLOAT", + { + "default": 8.0, + "min": 0.0, + "max": 100.0, + "step": 0.1, + "round": 0.01, + "tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt however too high values will negatively impact quality." + } + ], + "sampler_name": [ + [ + "euler", + "euler_cfg_pp", + "euler_ancestral", + "euler_ancestral_cfg_pp", + "heun", + "heunpp2", + "exp_heun_2_x0", + "exp_heun_2_x0_sde", + "dpm_2", + "dpm_2_ancestral", + "lms", + "dpm_fast", + "dpm_adaptive", + "dpmpp_2s_ancestral", + "dpmpp_2s_ancestral_cfg_pp", + "dpmpp_sde", + "dpmpp_sde_gpu", + "dpmpp_2m", + "dpmpp_2m_cfg_pp", + "dpmpp_2m_sde", + "dpmpp_2m_sde_gpu", + "dpmpp_2m_sde_heun", + "dpmpp_2m_sde_heun_gpu", + "dpmpp_3m_sde", + "dpmpp_3m_sde_gpu", + "ddpm", + "lcm", + "ipndm", + "ipndm_v", + "deis", + "res_multistep", + "res_multistep_cfg_pp", + "res_multistep_ancestral", + "res_multistep_ancestral_cfg_pp", + "gradient_estimation", + "gradient_estimation_cfg_pp", + "er_sde", + "seeds_2", + "seeds_3", + "sa_solver", + "sa_solver_pece", + "ddim", + "uni_pc", + "uni_pc_bh2" + ], + { + "tooltip": "The algorithm used when sampling, this can affect the quality, speed, and style of the generated output." + } + ], + "scheduler": [ + [ + "simple", + "sgm_uniform", + "karras", + "exponential", + "ddim_uniform", + "beta", + "normal", + "linear_quadratic", + "kl_optimal" + ], + { + "tooltip": "The scheduler controls how noise is gradually removed to form the image." + } + ], + "positive": [ + "CONDITIONING", + { + "tooltip": "The conditioning describing the attributes you want to include in the image." + } + ], + "negative": [ + "CONDITIONING", + { + "tooltip": "The conditioning describing the attributes you want to exclude from the image." + } + ], + "latent_image": [ + "LATENT", + { + "tooltip": "The latent image to denoise." + } + ], + "denoise": [ + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "The amount of denoising applied, lower values will maintain the structure of the initial image allowing for image to image sampling." + } + ] + } + }, + "input_order": { + "required": [ + "model", + "seed", + "steps", + "cfg", + "sampler_name", + "scheduler", + "positive", + "negative", + "latent_image", + "denoise" + ] + }, + "is_input_list": false, + "output": [ + "LATENT" + ], + "output_is_list": [ + false + ], + "output_name": [ + "LATENT" + ], + "name": "KSampler", + "display_name": "KSampler", + "description": "Uses the provided model, positive and negative conditioning to denoise the latent image.", + "python_module": "nodes", + "category": "sampling", + "output_node": false, + "has_intermediate_output": false, + "output_tooltips": [ + "The denoised latent." + ], + "search_aliases": [ + "sampler", + "sample", + "generate", + "denoise", + "diffuse", + "txt2img", + "img2img" + ] + }, + "SaveImage": { + "input": { + "required": { + "images": [ + "IMAGE", + { + "tooltip": "The images to save." + } + ], + "filename_prefix": [ + "STRING", + { + "default": "ComfyUI", + "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes." + } + ] + }, + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO" + } + }, + "input_order": { + "required": [ + "images", + "filename_prefix" + ], + "hidden": [ + "prompt", + "extra_pnginfo" + ] + }, + "is_input_list": false, + "output": [], + "output_is_list": [], + "output_name": [], + "name": "SaveImage", + "display_name": "Save Image", + "description": "Saves the input images to your ComfyUI output directory.", + "python_module": "nodes", + "category": "image", + "output_node": true, + "has_intermediate_output": false, + "search_aliases": [ + "save", + "save image", + "export image", + "output image", + "write image", + "download" + ], + "essentials_category": "Basics" + }, + "VAEDecode": { + "input": { + "required": { + "samples": [ + "LATENT", + { + "tooltip": "The latent to be decoded." + } + ], + "vae": [ + "VAE", + { + "tooltip": "The VAE model used for decoding the latent." + } + ] + } + }, + "input_order": { + "required": [ + "samples", + "vae" + ] + }, + "is_input_list": false, + "output": [ + "IMAGE" + ], + "output_is_list": [ + false + ], + "output_name": [ + "IMAGE" + ], + "name": "VAEDecode", + "display_name": "VAE Decode", + "description": "Decodes latent images back into pixel space images.", + "python_module": "nodes", + "category": "latent", + "output_node": false, + "has_intermediate_output": false, + "output_tooltips": [ + "The decoded image." + ], + "search_aliases": [ + "decode", + "decode latent", + "latent to image", + "render latent" + ] + } +} \ No newline at end of file diff --git a/comfy_cli/bench/fixtures/txt2img_seed.json b/comfy_cli/bench/fixtures/txt2img_seed.json new file mode 100644 index 00000000..4338b436 --- /dev/null +++ b/comfy_cli/bench/fixtures/txt2img_seed.json @@ -0,0 +1,547 @@ +{ + "id": "2ba0b800-2f13-4f21-b8d6-c6cdb0152cae", + "revision": 0, + "last_node_id": 16, + "last_link_id": 9, + "nodes": [ + { + "id": 4, + "type": "CheckpointLoaderSimple", + "pos": [ + 10, + 300 + ], + "size": [ + 320, + 154.65625 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "slot_index": 0, + "links": [ + 1 + ] + }, + { + "name": "CLIP", + "type": "CLIP", + "slot_index": 1, + "links": [ + 3, + 5 + ] + }, + { + "name": "VAE", + "type": "VAE", + "slot_index": 2, + "links": [ + 8 + ] + } + ], + "properties": { + "Node name for S&R": "CheckpointLoaderSimple", + "cnr_id": "comfy-core", + "ver": "0.3.65", + "models": [ + { + "name": "v1-5-pruned-emaonly-fp16.safetensors", + "url": "https://huggingface.co/Comfy-Org/stable-diffusion-v1-5-archive/resolve/main/v1-5-pruned-emaonly-fp16.safetensors?download=true", + "directory": "checkpoints" + } + ] + }, + "widgets_values": [ + "v1-5-pruned-emaonly-fp16.safetensors" + ] + }, + { + "id": 3, + "type": "KSampler", + "pos": [ + 920, + 170 + ], + "size": [ + 320, + 480 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 1 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 4 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 6 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 2 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "slot_index": 0, + "links": [ + 7 + ] + } + ], + "properties": { + "Node name for S&R": "KSampler", + "cnr_id": "comfy-core", + "ver": "0.3.65" + }, + "widgets_values": [ + 685468484323813, + "randomize", + 20, + 8, + "euler", + "normal", + 1 + ] + }, + { + "id": 8, + "type": "VAEDecode", + "pos": [ + 1000, + 710 + ], + "size": [ + 225, + 96 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": 8 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "slot_index": 0, + "links": [ + 9 + ] + } + ], + "properties": { + "Node name for S&R": "VAEDecode", + "cnr_id": "comfy-core", + "ver": "0.3.65" + }, + "widgets_values": [] + }, + { + "id": 9, + "type": "SaveImage", + "pos": [ + 1270, + 170 + ], + "size": [ + 470, + 560 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 9 + } + ], + "outputs": [], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.65" + }, + "widgets_values": [ + "SD1.5" + ] + }, + { + "id": 7, + "type": "CLIPTextEncode", + "pos": [ + 430, + 530 + ], + "size": [ + 420, + 170 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 5 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 6 + ] + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode", + "cnr_id": "comfy-core", + "ver": "0.3.65" + }, + "widgets_values": [ + "text, watermark" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 5, + "type": "EmptyLatentImage", + "pos": [ + 490, + 900 + ], + "size": [ + 320, + 168 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "slot_index": 0, + "links": [ + 2 + ] + } + ], + "properties": { + "Node name for S&R": "EmptyLatentImage", + "cnr_id": "comfy-core", + "ver": "0.3.65" + }, + "widgets_values": [ + 512, + 512, + 1 + ] + }, + { + "id": 6, + "type": "CLIPTextEncode", + "pos": [ + 420, + 220 + ], + "size": [ + 430, + 260 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 3 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 4 + ] + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode", + "cnr_id": "comfy-core", + "ver": "0.3.65" + }, + "widgets_values": [ + "beautiful scenery nature glass bottle landscape, purple galaxy bottle," + ], + "color": "#232", + "bgcolor": "#353" + }, + { + "id": 15, + "type": "MarkdownNote", + "pos": [ + 400, + -320 + ], + "size": [ + 470, + 430 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "Note: Prompt", + "properties": {}, + "widgets_values": [ + "**Prompts usually have two types:**\n\n* **Positive Prompt:** Tells the model what you *want* to see.\n* **Negative Prompt:** Tells the model what you *don\u2019t want* to see.\n\nThink of it as:
\n\ud83d\udc49 Positive = \u201cDo this\u201d
\n\ud83d\udc49 Negative = \u201cDon\u2019t do this\u201d\n\n\nDifferent models may interpret prompts differently.
\nSome prefer short, simple phrases; others respond well to detailed descriptions or styles.\nExperiment to see how each model reacts.\n\nAbout SD1.5:
\nStable Diffusion 1.5 is one of the most popular base models.\nIt works best with short, clear prompts and simple concepts, and it has a natural, realistic visual style.\n" + ], + "color": "#432", + "bgcolor": "#000" + }, + { + "id": 14, + "type": "MarkdownNote", + "pos": [ + 1270, + 780 + ], + "size": [ + 470, + 130 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "Note: Output", + "properties": {}, + "widgets_values": [ + "Image will auto-save to the `ComfyUI/output` folder. You can also right-click on the `Save Image` node and then use the menu to save the image locally." + ], + "color": "#432", + "bgcolor": "#000" + }, + { + "id": 13, + "type": "MarkdownNote", + "pos": [ + 460, + 1180 + ], + "size": [ + 330, + 163.953125 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "Note: Image size", + "properties": {}, + "widgets_values": [ + "Different models are trained based on datasets of different image sizes. This workflow is using Stable Diffusion 1.5, which is trained based on 512x512 datasets. So, it doesn't perform well on large image sizes." + ], + "color": "#432", + "bgcolor": "#000" + }, + { + "id": 11, + "type": "MarkdownNote", + "pos": [ + -470, + 160 + ], + "size": [ + 400, + 530.890625 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "Note: Model link", + "properties": {}, + "widgets_values": [ + "[Tutorial](https://docs.comfy.org/tutorials/basic/text-to-image)\n\n\n## Model link\n\n**checkpoints**\n\n- [v1-5-pruned-emaonly-fp16.safetensors](https://huggingface.co/Comfy-Org/stable-diffusion-v1-5-archive/resolve/main/v1-5-pruned-emaonly-fp16.safetensors?download=true)\n\n\nModel Storage Location\n\n```\n\ud83d\udcc2 ComfyUI/\n\u251c\u2500\u2500 \ud83d\udcc2 models/\n\u2502 \u2514\u2500\u2500 \ud83d\udcc2 checkpoints/\n\u2502 \u2514\u2500\u2500 v1-5-pruned-emaonly-fp16.safetensors\n```\n\n\n## Report issue\nIf you have any problems running this workflow, please report template-related issues via this link: [report the template issue here](https://github.com/Comfy-Org/workflow_templates/issues)" + ], + "color": "#432", + "bgcolor": "#000" + } + ], + "links": [ + [ + 1, + 4, + 0, + 3, + 0, + "MODEL" + ], + [ + 2, + 5, + 0, + 3, + 3, + "LATENT" + ], + [ + 3, + 4, + 1, + 6, + 0, + "CLIP" + ], + [ + 4, + 6, + 0, + 3, + 1, + "CONDITIONING" + ], + [ + 5, + 4, + 1, + 7, + 0, + "CLIP" + ], + [ + 6, + 7, + 0, + 3, + 2, + "CONDITIONING" + ], + [ + 7, + 3, + 0, + 8, + 0, + "LATENT" + ], + [ + 8, + 4, + 2, + 8, + 1, + "VAE" + ], + [ + 9, + 8, + 0, + 9, + 0, + "IMAGE" + ] + ], + "groups": [ + { + "id": 1, + "title": "Step 1 - Load model", + "bounding": [ + -40, + 130, + 420, + 470 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + }, + { + "id": 2, + "title": "Step 3 - Image size", + "bounding": [ + 400, + 800, + 480, + 310 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + }, + { + "id": 3, + "title": "Step 2 - Prompt", + "bounding": [ + 400, + 130, + 480, + 640 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + } + ], + "config": {}, + "extra": { + "ds": { + "scale": 0.5131581182307069, + "offset": [ + 979.5226642853634, + 273.924658465434 + ] + }, + "frontendVersion": "1.42.15", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/comfy_cli/bench/out/.gitignore b/comfy_cli/bench/out/.gitignore new file mode 100644 index 00000000..397002ef --- /dev/null +++ b/comfy_cli/bench/out/.gitignore @@ -0,0 +1,4 @@ +# Generated benchmark output — do not commit run artifacts. +*.ndjson +*_variants/ +!.gitkeep diff --git a/comfy_cli/bench/out/.gitkeep b/comfy_cli/bench/out/.gitkeep new file mode 100644 index 00000000..8ace8e1a --- /dev/null +++ b/comfy_cli/bench/out/.gitkeep @@ -0,0 +1 @@ +# Keep the out/ dir in git; generated NDJSON is ignored (see .gitignore). diff --git a/comfy_cli/bench/run_arm_b.py b/comfy_cli/bench/run_arm_b.py new file mode 100644 index 00000000..8dd232e4 --- /dev/null +++ b/comfy_cli/bench/run_arm_b.py @@ -0,0 +1,926 @@ +"""Arm-B runner for the BE-2302 A/B micro-edit benchmark — a **SANCTIONED PROXY**. + +Arm B of BE-2302 is defined as "4-5 generic micro-edit tools operating on a temp +JSON workflow on disk." This module is a minimal agent loop (Claude Opus +``claude-opus-4-8`` over the raw Anthropic Messages API) that exposes exactly the +``comfy workflow`` micro-edit substrate already in this repo as tools: + + read → ``slots`` (comfy workflow slots) + ``cql`` (node/type lookup) + edit → ``set_slot`` (comfy workflow set-slot) + produce-variant→ ``vary`` (comfy workflow vary) + +It runs the SAME t1–t4 task prompts as arm A (vendored byte-identical from child 1's +``tasks.mjs`` into ``bench/tasks.json``) and emits per-(task, turn) NDJSON rows shaped +IDENTICALLY to arm A's ``arm-a.ndjson`` so child 1's ``report.mjs`` consumes both +directly. Cost is recomputed from the Anthropic ``usage`` block with the SAME pricing +table as ``comfy-inapp-agent/agent-server/usage.mjs`` (Opus 4.8: $5/M in, $25/M out, +cache read 0.1×, cache write 1.25×) so the two arms are dollar-comparable. + +──────────────────────────────────────────────────────────────────────────────── +PROXY CAVEAT — this is NOT Kishore's real CLI-agent prototype. +──────────────────────────────────────────────────────────────────────────────── +BE-2302 sanctions this as a *proxy* for a not-yet-accessible CLI agent prototype. +Every emitted row is stamped ``"proxy": true`` and ``"arm": "B"`` so no downstream +reader mistakes it for the real thing. The swap-in seam is deliberately narrow: + + * The MODEL driver is ``Driver`` (see ``run_task``). To swap in the real + prototype, implement a driver that speaks the prototype's protocol but keeps + calling ``ToolDispatcher`` (the comfy-cli substrate) and emitting rows through + ``build_row`` — the telemetry + NDJSON shape then stay identical for free. + * The MODEL CLIENT is injected (``client`` arg). ``StubClient`` (offline, + deterministic) drives ``--dry-run`` and the unit tests; ``build_live_client`` + returns a real ``anthropic.Anthropic`` for ``--live``. Point the client at the + prototype's endpoint, or replace ``Driver`` wholesale, and nothing else moves. + +See ``bench/README.md`` for the live-run recipe and the full swap-in contract. +""" + +from __future__ import annotations + +import argparse +import json +import shutil +import sys +import time +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +# --------------------------------------------------------------------------- +# Pricing — mirror of comfy-inapp-agent/agent-server/usage.mjs PRICING so the two +# arms price a turn identically. Values are USD per million tokens. +# --------------------------------------------------------------------------- + +PRICING: list[tuple[str, float, float]] = [ + # (id-prefix, input $/MTok, output $/MTok) + ("claude-fable-5", 10.0, 50.0), + ("claude-opus-4-8", 5.0, 25.0), + ("claude-opus-4-7", 5.0, 25.0), + ("claude-opus-4-6", 5.0, 25.0), + ("claude-sonnet-4-6", 3.0, 15.0), + ("claude-haiku-4-5", 1.0, 5.0), +] +CACHE_READ_MULT = 0.1 +CACHE_WRITE_MULT = 1.25 + +DEFAULT_MODEL = "claude-opus-4-8" + +# Token fields the runner captures from the Anthropic ``usage`` block, emitted +# verbatim so report.mjs can recompute cost from them. +TOKEN_FIELDS = ( + "input_tokens", + "output_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", +) + + +def _round6(v: float) -> float: + return round(v * 1e6) / 1e6 + + +def resolve_pricing(model: str | None) -> tuple[str, float, float] | None: + """Longest-prefix match of a model id against the pricing table; None when unknown.""" + if not isinstance(model, str) or not model: + return None + best: tuple[str, float, float] | None = None + for entry in PRICING: + prefix = entry[0] + if model.startswith(prefix) and (best is None or len(prefix) > len(best[0])): + best = entry + return best + + +def estimate_cost_usd(model: str | None, usage: dict[str, Any]) -> float | None: + """USD cost of one turn's tokens at list pricing. None when the model is unknown. + + Mirrors ``usage.mjs`` ``estimateCostUsd`` exactly (cache read 0.1×, write 1.25×). + """ + p = resolve_pricing(model) + if p is None: + return None + _, in_rate, out_rate = p + n = lambda v: v if isinstance(v, int | float) and not isinstance(v, bool) else 0 # noqa: E731 + return _round6( + ( + n(usage.get("input_tokens")) * in_rate + + n(usage.get("output_tokens")) * out_rate + + n(usage.get("cache_read_input_tokens")) * in_rate * CACHE_READ_MULT + + n(usage.get("cache_creation_input_tokens")) * in_rate * CACHE_WRITE_MULT + ) + / 1e6 + ) + + +def parse_usage(usage: Any) -> dict[str, int]: + """Normalize an Anthropic ``usage`` block (SDK object or dict) → the 4 token counts. + + Missing / non-numeric fields default to 0 so a partial usage block never crashes + the loop. This is the unit-tested seam between the API response and the NDJSON row. + """ + + def get(name: str) -> int: + if usage is None: + return 0 + v = usage.get(name) if isinstance(usage, dict) else getattr(usage, name, None) + if isinstance(v, bool) or not isinstance(v, int | float): + return 0 + return int(v) + + return {f: get(f) for f in TOKEN_FIELDS} + + +# --------------------------------------------------------------------------- +# Tool schemas — the comfy-cli micro-edit substrate exposed to the model. +# --------------------------------------------------------------------------- + +TOOLS: list[dict[str, Any]] = [ + { + "name": "slots", + "description": ( + "READ. List the tweakable slots of the workflow currently on disk " + "(`comfy workflow slots`). Each slot is {address, type, current_value}; " + "address is `.`. Call this first to discover what you can edit." + ), + "input_schema": {"type": "object", "properties": {}, "additionalProperties": False}, + }, + { + "name": "cql", + "description": ( + "READ. Validate a node type / inspect the object_info catalog (the CQL layer). " + "Pass `node_type` to fetch that node class's input/output schema, or omit it to " + "list every node class available offline. Use this to check a node exists before editing." + ), + "input_schema": { + "type": "object", + "properties": {"node_type": {"type": "string", "description": "Node class to look up, e.g. 'KSampler'."}}, + "additionalProperties": False, + }, + }, + { + "name": "set_slot", + "description": ( + "EDIT. Apply one or more slot overrides to the workflow in place " + "(`comfy workflow set-slot`). `overrides` maps slot address → new value, " + "e.g. {'6.text': 'a cat', '3.steps': 25}." + ), + "input_schema": { + "type": "object", + "properties": {"overrides": {"type": "object", "description": "address → value map."}}, + "required": ["overrides"], + "additionalProperties": False, + }, + }, + { + "name": "vary", + "description": ( + "PRODUCE-VARIANTS. Expand the workflow into N variants from a per-slot value list " + "(`comfy workflow vary`). `slots` maps address → list of values; all lists must be the " + "same length N. Returns the variant count (variants are written to a temp dir, not to the model)." + ), + "input_schema": { + "type": "object", + "properties": {"slots": {"type": "object", "description": "address → list-of-values map."}}, + "required": ["slots"], + "additionalProperties": False, + }, + }, +] + +SYSTEM_PROMPT = ( + "You are a proxy for a ComfyUI CLI micro-edit agent. You edit a frontend-format workflow JSON " + "that lives on disk via four tools (slots, cql, set_slot, vary) — you NEVER receive or emit the " + "full workflow JSON, only compact slot manifests. Discover slots with `slots`, validate node " + "types with `cql`, make edits with `set_slot`, and produce variants with `vary`. Your toolset can " + "only EDIT an existing graph and its known node catalog; it cannot add arbitrary new node types. " + "If a task needs node classes that `cql` reports are unavailable, say so plainly and stop." +) + + +# --------------------------------------------------------------------------- +# Tool dispatch — every call hits the real comfy-cli CQL engine against the +# committed object_info fixture + the temp workflow on disk. +# --------------------------------------------------------------------------- + + +class ToolDispatcher: + """Maps a model tool call onto the actual comfy-cli micro-edit primitives. + + The graph is built once from the offline ``object_info`` fixture; the workflow + is (re)read from ``workflow_path`` on every call so a ``set_slot`` write is + visible to the next ``slots`` read — a genuine on-disk round-trip, no full-JSON + round-trip to the model. + """ + + def __init__(self, object_info_path: str | Path, workflow_path: str | Path, variant_dir: str | Path): + from comfy_cli.cql.engine import Graph + + self.graph = Graph.load(input_path=str(object_info_path)) + self.workflow_path = Path(workflow_path) + self.variant_dir = Path(variant_dir) + + def _load_workflow(self) -> dict[str, Any]: + return json.loads(self.workflow_path.read_text(encoding="utf-8")) + + def dispatch(self, name: str, tool_input: dict[str, Any]) -> tuple[dict[str, Any], bool]: + """Run one tool. Returns (result_dict, is_error). Never raises.""" + try: + handler = { + "slots": self._slots, + "cql": self._cql, + "set_slot": self._set_slot, + "vary": self._vary, + }.get(name) + if handler is None: + return {"error": f"unknown tool {name!r}"}, True + return handler(tool_input or {}) + except Exception as e: # dispatch must be total — surface the error to the model + return {"error": f"{type(e).__name__}: {e}"}, True + + # -- read tools -- + + def _slots(self, _inp: dict[str, Any]) -> tuple[dict[str, Any], bool]: + wf = self._load_workflow() + schema = self.graph.get_template_schema(self.workflow_path.stem, wf) + slots = [ + {"address": s.get("address"), "type": s.get("type"), "current_value": s.get("current_value")} + for s in schema.get("slots") or [] + ] + return {"count": len(slots), "slots": slots}, False + + def _cql(self, inp: dict[str, Any]) -> tuple[dict[str, Any], bool]: + node_type = inp.get("node_type") + if not node_type: + return {"available_node_types": [m.id for m in self.graph.all_nodes()]}, False + m = self.graph.node(node_type) + if m is None: + import difflib + + names = [n.id for n in self.graph.all_nodes()] + close = difflib.get_close_matches(node_type, names, n=5, cutoff=0.4) + return {"node_type": node_type, "found": False, "suggestions": close}, False + return {"node_type": node_type, "found": True, "schema": self.graph.morphism_to_dict(m)}, False + + # -- edit tool -- + + def _set_slot(self, inp: dict[str, Any]) -> tuple[dict[str, Any], bool]: + overrides = inp.get("overrides") + if not isinstance(overrides, dict) or not overrides: + return {"error": "set_slot requires a non-empty `overrides` object (address → value)."}, True + wf = self._load_workflow() + try: + new_wf, warnings = self.graph.apply_slots(wf, overrides) + except ValueError as e: + return {"error": str(e), "hint": "run `slots` to see valid addresses + types"}, True + self.workflow_path.write_text(json.dumps(new_wf, indent=2), encoding="utf-8") + return {"applied": list(overrides.keys()), "warnings": warnings}, False + + # -- produce-variants tool -- + + def _vary(self, inp: dict[str, Any]) -> tuple[dict[str, Any], bool]: + by_addr = inp.get("slots") + if not isinstance(by_addr, dict) or not by_addr: + return {"error": "vary requires a non-empty `slots` object (address → list-of-values)."}, True + lengths = {a: len(v) for a, v in by_addr.items() if isinstance(v, list)} + if len(lengths) != len(by_addr): + return {"error": "every `slots` value must be a JSON array."}, True + if len(set(lengths.values())) != 1: + return {"error": f"all `slots` lists must be the same length; got {lengths}."}, True + n = next(iter(lengths.values())) + variations = [{a: v[i] for a, v in by_addr.items()} for i in range(n)] + wf = self._load_workflow() + try: + workflows, warnings = self.graph.expand_variations(wf, variations) + except ValueError as e: + return {"error": str(e)}, True + self.variant_dir.mkdir(parents=True, exist_ok=True) + written = [] + for i, w in enumerate(workflows): + target = self.variant_dir / f"{self.workflow_path.stem}_{i:03d}.json" + target.write_text(json.dumps(w, indent=2), encoding="utf-8") + written.append(target.name) + return {"count": len(workflows), "written": written, "warnings": warnings}, False + + +# --------------------------------------------------------------------------- +# Model-response normalization — one interface over the real SDK objects and the +# stub's plain objects so ``run_task`` never branches on client type. +# --------------------------------------------------------------------------- + + +def _battr(block: Any, name: str, default: Any = None) -> Any: + if isinstance(block, dict): + return block.get(name, default) + return getattr(block, name, default) + + +@dataclass +class TurnTelemetry: + """Accumulates one turn's telemetry across its (possibly many) API round-trips.""" + + input_tokens: int = 0 + output_tokens: int = 0 + cache_read_input_tokens: int = 0 + cache_creation_input_tokens: int = 0 + num_api_calls: int = 0 + tool_call_count: int = 0 + tool_input_bytes: list[int] = field(default_factory=list) + tool_output_bytes: list[int] = field(default_factory=list) + duration_ms: float = 0.0 + + def add_usage(self, usage: dict[str, int]) -> None: + for f in TOKEN_FIELDS: + setattr(self, f, getattr(self, f) + usage[f]) + self.num_api_calls += 1 + + def add_tool(self, input_bytes: int, output_bytes: int) -> None: + self.tool_call_count += 1 + self.tool_input_bytes.append(input_bytes) + self.tool_output_bytes.append(output_bytes) + + +def build_row( + *, + task: dict[str, Any], + turn: int, + model: str, + tel: TurnTelemetry, + outcome: str = "ok", + note: str | None = None, +) -> dict[str, Any]: + """Assemble one NDJSON row, shaped identically to arm A's ``arm-a.ndjson``. + + Extra keys (``proxy``, ``outcome``, ``note``) are additive — report.mjs reads only + the shared schema, so the proxy labelling rides along harmlessly. + """ + token_usage = {f: getattr(tel, f) for f in TOKEN_FIELDS} + cost = estimate_cost_usd(model, token_usage) + cost = 0.0 if cost is None else cost + ti, to = tel.tool_input_bytes, tel.tool_output_bytes + payload = [i + o for i, o in zip(ti, to)] + duration_ms = int(round(tel.duration_ms)) + return { + "arm": "B", + "proxy": True, # SANCTIONED PROXY — not Kishore's real prototype (see module docstring) + "task": task["id"], + "title": task["title"], + "turn": turn, + "msg_id": f"bench-{task['id']}-{turn}", + "model": model, + **token_usage, + "cost_usd": cost, + # No separate SDK cost figure on the raw Messages API — mirror our own estimate. + "sdk_cost_usd": cost, + "num_turns": tel.num_api_calls, + "duration_ms": duration_ms, + "duration_api_ms": duration_ms, + "wall_ms": duration_ms, + "tool_call_count": tel.tool_call_count, + "tool_input_bytes_total": sum(ti), + "tool_input_bytes_max": max(ti) if ti else 0, + "tool_output_bytes_total": sum(to), + "tool_output_bytes_max": max(to) if to else 0, + "tool_payload_bytes_total": sum(payload), + "tool_payload_bytes_max": max(payload) if payload else 0, + # Arm B never compacts a full transcript (it round-trips slot manifests, not JSON). + "compactions": 0, + "outcome": outcome, + "note": note, + } + + +# --------------------------------------------------------------------------- +# Driver — the agent loop. This is the swap-in seam for the real prototype. +# --------------------------------------------------------------------------- + +MAX_API_CALLS_PER_TURN = 12 # guard against a tool-loop that never ends + + +class Driver: + """Runs one task's prompts through the injected model ``client`` + ``dispatcher``. + + To swap in Kishore's real CLI-agent prototype: subclass/replace this driver with + one that speaks the prototype's protocol, keep ``dispatcher`` (the comfy-cli + substrate) and ``build_row`` (the telemetry shape), and everything downstream — + NDJSON, report.mjs — stays identical. + """ + + def __init__(self, client: Any, model: str = DEFAULT_MODEL, max_tokens: int = 2048): + self.client = client + self.model = model + self.max_tokens = max_tokens + + def _create(self, messages: list[dict[str, Any]]) -> Any: + return self.client.messages.create( + model=self.model, + max_tokens=self.max_tokens, + system=SYSTEM_PROMPT, + tools=TOOLS, + messages=messages, + ) + + def run_turn( + self, dispatcher: ToolDispatcher, prompt: str, messages: list[dict[str, Any]] | None = None + ) -> TurnTelemetry: + """Drive one user prompt to completion, dispatching tools; return its telemetry. + + ``messages`` is the running transcript. Pass the SAME list across a task's turns + (t4's one-session, three-edit flow) so the transcript grows turn-over-turn and the + cache/regrowth telemetry mirrors arm A's resumed session; omit it for a fresh, + single-turn task and each turn starts clean. + """ + if messages is None: + messages = [] + tel = TurnTelemetry() + messages.append({"role": "user", "content": prompt}) + for _ in range(MAX_API_CALLS_PER_TURN): + t0 = time.perf_counter() + resp = self._create(messages) + # A stub can supply a synthetic latency so --dry-run timings are + # deterministic; live runs measure the real wall time. + stub_latency = _battr(resp, "stub_latency_ms") + tel.duration_ms += stub_latency if stub_latency is not None else (time.perf_counter() - t0) * 1000.0 + tel.add_usage(parse_usage(_battr(resp, "usage"))) + + content = _battr(resp, "content") or [] + messages.append({"role": "assistant", "content": content}) + tool_uses = [b for b in content if _battr(b, "type") == "tool_use"] + if _battr(resp, "stop_reason") != "tool_use" or not tool_uses: + break + + tool_results = [] + for tu in tool_uses: + tu_input = _battr(tu, "input") or {} + result, is_error = dispatcher.dispatch(_battr(tu, "name"), tu_input) + result_json = json.dumps(result) + tel.add_tool(len(json.dumps(tu_input)), len(result_json)) + tool_results.append( + { + "type": "tool_result", + "tool_use_id": _battr(tu, "id"), + "content": result_json, + "is_error": is_error, + } + ) + messages.append({"role": "user", "content": tool_results}) + return tel + + +def _seed_for_task(task: dict[str, Any], fixtures_dir: Path) -> Path: + """Resolve the on-disk seed graph a task edits. + + t2/t4 name a seed (``txt2img-seed`` / ``edit-session-seed``) → committed fixture. + t1/t3 build from scratch; arm B (a micro-edit toolset) has no from-scratch node + creation, so it seeds from the txt2img template and micro-edits it — the honest + proxy of "build small". t3's true ceiling is recorded as a RESULT (see run_arm_b). + """ + name = task.get("seedWorkflow") + mapping = {"txt2img-seed": "txt2img_seed.json", "edit-session-seed": "edit_session_seed.json"} + return fixtures_dir / mapping.get(name, "txt2img_seed.json") + + +# t3 is the documented ceiling: arm B's slots/set-slot/vary can only EDIT an existing +# graph, and the committed object_info covers only the 6 base SD1.5 classes (no +# LoRA/ControlNet/refiner/upscaler/face-detailer), so a ~150-node BUILD is out of reach. +_T3_CEILING_NOTE = ( + "arm B micro-edit toolset cannot build a ~150-node graph; object_info fixture " + "covers only the 6 base SD1.5 node classes (no LoRA/ControlNet/refiner/upscaler). " + "Recorded as a RESULT per BE-2309." +) + + +def _task_outcome(task_id: str) -> tuple[str, str | None]: + """Per-task (outcome, note) stamp — 'ceiling' RESULT for t3, plain 'ok' otherwise.""" + if task_id == "t3": + return "ceiling", _T3_CEILING_NOTE + return "ok", None + + +def run_task( + task: dict[str, Any], + *, + driver: Driver, + fixtures_dir: Path, + object_info_path: Path, + work_dir: Path, + on_row: Callable[[dict[str, Any]], None], +) -> None: + """Run every turn of a task on a FRESH temp copy of its seed, emitting one row per turn. + + A task's turns share ONE transcript (``messages``) so a multi-prompt task (t4) grows its + context turn-over-turn — the same one-session semantics arm A measures. + """ + seed = _seed_for_task(task, fixtures_dir) + workflow_path = work_dir / f"{task['id']}_workflow.json" + shutil.copyfile(seed, workflow_path) + dispatcher = ToolDispatcher(object_info_path, workflow_path, work_dir / f"{task['id']}_variants") + + messages: list[dict[str, Any]] = [] + for i, prompt in enumerate(task["prompts"], start=1): + tel = driver.run_turn(dispatcher, prompt, messages) + outcome, note = _task_outcome(task["id"]) + on_row(build_row(task=task, turn=i, model=driver.model, tel=tel, outcome=outcome, note=note)) + + +# --------------------------------------------------------------------------- +# Stub client — offline, deterministic. Drives --dry-run and the unit tests. +# --------------------------------------------------------------------------- + + +class _StubBlock: + """A minimal content block matching the attribute surface run_task reads.""" + + def __init__(self, type: str, **kw: Any): + self.type = type + self.text = kw.get("text") + self.name = kw.get("name") + self.input = kw.get("input") + self.id = kw.get("id") + + +class _StubResponse: + def __init__(self, content: list[_StubBlock], stop_reason: str, usage: dict[str, int], latency_ms: float): + self.content = content + self.stop_reason = stop_reason + self.usage = usage + self.stub_latency_ms = latency_ms + + +def _text(t: str) -> _StubBlock: + return _StubBlock("text", text=t) + + +def _tool(name: str, tid: str, tool_input: dict[str, Any]) -> _StubBlock: + return _StubBlock("tool_use", name=name, id=tid, input=tool_input) + + +# A scripted proxy plan per (task, turn): the ordered list of responses the stub +# returns. Each response is (content_blocks, stop_reason, usage). Addresses are the +# real slot addresses of the committed sd15 seed fixtures, so the tools genuinely +# read/edit the temp workflow on disk. Token counts are synthetic-but-plausible and +# show modest cache regrowth across a session (t4). +def _usage(inp: int, out: int, cr: int, cw: int) -> dict[str, int]: + return { + "input_tokens": inp, + "output_tokens": out, + "cache_read_input_tokens": cr, + "cache_creation_input_tokens": cw, + } + + +def default_stub_plan() -> dict[tuple[str, int], list[tuple[list[_StubBlock], str, dict[str, int]]]]: + """Deterministic proxy scripts keyed by (task_id, 1-based turn index).""" + return { + # t1 — "build small": discover slots, set the requested params, produce 2 variants. + ("t1", 1): [ + ([_tool("slots", "t1a", {})], "tool_use", _usage(700, 120, 2000, 1800)), + ( + [ + _tool( + "set_slot", + "t1b", + { + "overrides": { + "6.text": "a golden retriever puppy playing in autumn leaves, cinematic lighting", + "3.steps": 25, + "3.cfg": 7, + "3.sampler_name": "euler", + "3.scheduler": "normal", + } + }, + ) + ], + "tool_use", + _usage(950, 240, 4200, 900), + ), + ( + [_tool("vary", "t1c", {"slots": {"3.seed": [1, 2]}})], + "tool_use", + _usage(1100, 90, 5600, 300), + ), + ( + [_text("Built the SDXL txt2img template and produced 2 seed variants.")], + "end_turn", + _usage(1200, 160, 6400, 0), + ), + ], + # t2 — micro-edit ONE parameter (the canonical arm-B case). + ("t2", 1): [ + ([_tool("slots", "t2a", {})], "tool_use", _usage(650, 80, 1800, 1600)), + ( + [ + _tool( + "set_slot", + "t2b", + {"overrides": {"6.text": "a photo of a cat sleeping in a sunbeam on a wooden floor"}}, + ) + ], + "tool_use", + _usage(820, 130, 3400, 200), + ), + ( + [_text("Changed only the positive prompt; every other node is untouched.")], + "end_turn", + _usage(900, 110, 3900, 0), + ), + ], + # t3 — the CEILING: probe the catalog for the classes the task needs, find + # them absent, and stop. (run_task stamps outcome="ceiling".) + ("t3", 1): [ + ([_tool("cql", "t3a", {"node_type": "LoraLoader"})], "tool_use", _usage(900, 140, 2200, 2000)), + ([_tool("cql", "t3b", {"node_type": "ControlNetLoader"})], "tool_use", _usage(1050, 120, 3600, 400)), + ( + [ + _text( + "The requested LoRA / ControlNet / refiner / upscaler classes are not in the offline " + "object_info catalog, and my micro-edit tools cannot create ~150 new nodes. Stopping — " + "this is the documented arm-B ceiling." + ) + ], + "end_turn", + _usage(1200, 260, 4300, 0), + ), + ], + # t4 — three successive edits + read-backs in ONE session (transcript regrowth). + ("t4", 1): [ + ( + [_tool("set_slot", "t4a", {"overrides": {"3.steps": 40, "3.sampler_name": "euler"}})], + "tool_use", + _usage(700, 150, 2400, 2200), + ), + ([_tool("slots", "t4b", {})], "tool_use", _usage(880, 90, 4100, 500)), + ( + [_text("Set KSampler to 40 steps and confirmed the sampler settings on read-back.")], + "end_turn", + _usage(1000, 130, 4800, 0), + ), + ], + ("t4", 2): [ + ( + [ + _tool( + "set_slot", + "t4c", + { + "overrides": { + "6.text": "a dense bioluminescent alien jungle at night, volumetric fog, highly detailed" + } + }, + ) + ], + "tool_use", + _usage(760, 160, 6200, 700), + ), + ([_tool("slots", "t4d", {})], "tool_use", _usage(1020, 100, 7800, 300)), + ( + [_text("Updated the positive prompt and confirmed the change landed on read-back.")], + "end_turn", + _usage(1150, 120, 8600, 0), + ), + ], + ("t4", 3): [ + ( + [_tool("set_slot", "t4e", {"overrides": {"9.filename_prefix": "jungle_alt"}})], + "tool_use", + _usage(820, 170, 9400, 600), + ), + ([_tool("slots", "t4f", {})], "tool_use", _usage(1120, 110, 11000, 300)), + ( + [_text("Set the SaveImage filename prefix to jungle_alt and summarized the final graph.")], + "end_turn", + _usage(1260, 200, 11800, 0), + ), + ], + } + + +class StubClient: + """Offline, scripted Anthropic-compatible client. ``messages.create`` pops the + next response from the active (task, turn) script set by ``begin_turn``.""" + + def __init__(self, plan: dict[tuple[str, int], list[tuple[list[_StubBlock], str, dict[str, int]]]] | None = None): + self.plan = plan if plan is not None else default_stub_plan() + self._queue: list[tuple[list[_StubBlock], str, dict[str, int]]] = [] + self.messages = self._Messages(self) + + def begin_turn(self, task_id: str, turn: int) -> None: + self._queue = list(self.plan.get((task_id, turn), [])) + if not self._queue: + # A turn with no script terminates immediately (single empty end_turn). + self._queue = [([_text("(no scripted action)")], "end_turn", _usage(300, 40, 0, 0))] + + class _Messages: + def __init__(self, outer: StubClient): + self._outer = outer + + def create(self, **_kwargs: Any) -> _StubResponse: + if not self._outer._queue: + # Ran past the script (would only happen if a turn tool-loops longer + # than planned) — end cleanly rather than hang. + return _StubResponse([_text("(script exhausted)")], "end_turn", _usage(200, 20, 0, 0), 5.0) + content, stop, usage = self._outer._queue.pop(0) + return _StubResponse(content, stop, usage, latency_ms=250.0) + + +class _StubDriver(Driver): + """Driver bound to a StubClient — announces each turn to the stub so it serves + the right (task, turn) script.""" + + def run_turn_for( + self, + dispatcher: ToolDispatcher, + task_id: str, + turn: int, + prompt: str, + messages: list[dict[str, Any]] | None = None, + ) -> TurnTelemetry: + self.client.begin_turn(task_id, turn) + return self.run_turn(dispatcher, prompt, messages) + + +def run_task_stub( + task: dict[str, Any], + *, + driver: _StubDriver, + fixtures_dir: Path, + object_info_path: Path, + work_dir: Path, + on_row: Callable[[dict[str, Any]], None], +) -> None: + """run_task variant that tells the stub which (task, turn) script to serve. + + Mirrors ``run_task``: one shared transcript across a task's turns so t4 grows in-session. + """ + seed = _seed_for_task(task, fixtures_dir) + workflow_path = work_dir / f"{task['id']}_workflow.json" + shutil.copyfile(seed, workflow_path) + dispatcher = ToolDispatcher(object_info_path, workflow_path, work_dir / f"{task['id']}_variants") + messages: list[dict[str, Any]] = [] + for i, prompt in enumerate(task["prompts"], start=1): + tel = driver.run_turn_for(dispatcher, task["id"], i, prompt, messages) + outcome, note = _task_outcome(task["id"]) + on_row(build_row(task=task, turn=i, model=driver.model, tel=tel, outcome=outcome, note=note)) + + +# --------------------------------------------------------------------------- +# Live client — real Anthropic SDK. Imported lazily so --dry-run + tests stay offline. +# --------------------------------------------------------------------------- + + +def build_live_client() -> Any: + """Return a real ``anthropic.Anthropic`` client (reads ANTHROPIC_API_KEY from env). + + Swap-in seam: to point arm B at Kishore's real CLI-agent prototype, return a + client object exposing ``messages.create(...)`` with the same response surface + (``.content``, ``.stop_reason``, ``.usage``), or replace ``Driver`` entirely. + """ + try: + import anthropic + except ImportError as e: # pragma: no cover - only hit on a live run without the extra + raise SystemExit( + "The `anthropic` package is required for a live run. Install the bench extra:\n" + " pip install -e '.[bench]'\n" + "and set ANTHROPIC_API_KEY. Use --dry-run for an offline stubbed run." + ) from e + return anthropic.Anthropic() + + +# --------------------------------------------------------------------------- +# Data locations + task loading +# --------------------------------------------------------------------------- + +_HERE = Path(__file__).resolve().parent +FIXTURES_DIR = _HERE / "fixtures" +TASKS_PATH = _HERE / "tasks.json" +DEFAULT_OBJECT_INFO = FIXTURES_DIR / "object_info.json" +DEFAULT_OUT = _HERE / "out" / "arm-b.ndjson" + + +def load_tasks(path: str | Path = TASKS_PATH) -> list[dict[str, Any]]: + data = json.loads(Path(path).read_text(encoding="utf-8")) + tasks = data["tasks"] + # Normalize: every task has a `prompts` list (matching child 1's taskPrompts()). + for t in tasks: + if "prompts" not in t or not isinstance(t["prompts"], list): + t["prompts"] = [t["prompt"]] if t.get("prompt") else [] + return tasks + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def run( + *, + dry_run: bool, + out_path: Path, + object_info_path: Path, + tasks_path: Path, + model: str, + work_dir: Path, + only: list[str] | None = None, +) -> list[dict[str, Any]]: + """Run the whole matrix, write NDJSON to ``out_path``, and return the rows.""" + tasks = load_tasks(tasks_path) + if only: + tasks = [t for t in tasks if t["id"] in set(only)] + work_dir.mkdir(parents=True, exist_ok=True) + rows: list[dict[str, Any]] = [] + out_path.parent.mkdir(parents=True, exist_ok=True) + + with out_path.open("w", encoding="utf-8") as fh: + + def on_row(row: dict[str, Any]) -> None: + rows.append(row) + fh.write(json.dumps(row) + "\n") + fh.flush() + + if dry_run: + driver = _StubDriver(StubClient(), model=model) + for task in tasks: + run_task_stub( + task, + driver=driver, + fixtures_dir=FIXTURES_DIR, + object_info_path=object_info_path, + work_dir=work_dir, + on_row=on_row, + ) + else: + driver = Driver(build_live_client(), model=model) + for task in tasks: + run_task( + task, + driver=driver, + fixtures_dir=FIXTURES_DIR, + object_info_path=object_info_path, + work_dir=work_dir, + on_row=on_row, + ) + return rows + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + prog="python -m comfy_cli.bench.run_arm_b", + description="Arm-B (CLI micro-edit PROXY) runner for the BE-2302 A/B benchmark.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Offline: drive a stubbed model against the committed object_info fixture (no API key, no network).", + ) + parser.add_argument( + "--live", action="store_true", help="Real run against the Anthropic API (needs ANTHROPIC_API_KEY)." + ) + parser.add_argument("--out", type=Path, default=DEFAULT_OUT, help=f"NDJSON output path (default: {DEFAULT_OUT}).") + parser.add_argument( + "--object-info", type=Path, default=DEFAULT_OBJECT_INFO, help="object_info JSON (offline catalog)." + ) + parser.add_argument("--tasks", type=Path, default=TASKS_PATH, help="tasks.json (shared t1–t4 prompts).") + parser.add_argument("--model", default=DEFAULT_MODEL, help=f"Model id (default: {DEFAULT_MODEL}).") + parser.add_argument( + "--work-dir", type=Path, default=None, help="Scratch dir for temp workflows (default: a tempdir)." + ) + parser.add_argument("--only", nargs="*", default=None, help="Restrict to specific task ids, e.g. --only t2 t4.") + args = parser.parse_args(argv) + + if not args.dry_run and not args.live: + parser.error("choose one of --dry-run (offline, stubbed) or --live (real Anthropic API).") + if args.dry_run and args.live: + parser.error("--dry-run and --live are mutually exclusive.") + + import tempfile + + work_dir = args.work_dir + tmp_ctx = None + if work_dir is None: + tmp_ctx = tempfile.TemporaryDirectory(prefix="arm-b-bench-") + work_dir = Path(tmp_ctx.name) + try: + rows = run( + dry_run=args.dry_run, + out_path=args.out, + object_info_path=args.object_info, + tasks_path=args.tasks, + model=args.model, + work_dir=Path(work_dir), + only=args.only, + ) + finally: + if tmp_ctx is not None: + tmp_ctx.cleanup() + + mode = "dry-run (stubbed, PROXY)" if args.dry_run else "live (PROXY)" + total_cost = sum(r["cost_usd"] for r in rows) + sys.stderr.write( + f"✓ arm B {mode}: wrote {len(rows)} row(s) → {args.out} (est ${total_cost:.4f}). " + f"Feed to child 1's report.mjs: `node bench/report.mjs --arm-b {args.out}`.\n" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/comfy_cli/bench/tasks.json b/comfy_cli/bench/tasks.json new file mode 100644 index 00000000..0ae18773 --- /dev/null +++ b/comfy_cli/bench/tasks.json @@ -0,0 +1,39 @@ +{ + "_comment": "Vendored byte-identical from comfy-inapp-agent/agent-server/bench/tasks.mjs (child 1, BE-2308). Prompt strings MUST stay in sync across both repos; see comfy_cli/bench/README.md.", + "tasks": [ + { + "id": "t1", + "title": "Build a small template workflow (~7–15 nodes) from scratch", + "seedWorkflow": null, + "prompts": [ + "Build a standard SDXL text-to-image workflow on a new canvas tab: a checkpoint loader, positive and negative CLIP text-encode prompts, an empty latent image, a KSampler, a VAE decode, and a save-image node, all wired together. Use the positive prompt \"a golden retriever puppy playing in autumn leaves, cinematic lighting\". Set the KSampler to 25 steps, cfg 7, the euler sampler and the normal scheduler. Build it in one batched send_graph_mutation into a new tab so I can revert in one click." + ] + }, + { + "id": "t2", + "title": "Micro-edit one parameter in a seed workflow", + "seedWorkflow": "txt2img-seed", + "prompts": [ + "The workflow on my canvas has a positive prompt that says the cat is on a windowsill. Change ONLY the positive prompt text to \"a photo of a cat sleeping in a sunbeam on a wooden floor\". Leave every other node and parameter exactly as it is." + ] + }, + { + "id": "t3", + "title": "Build / extend a large workflow (~150 nodes)", + "seedWorkflow": null, + "prompts": [ + "I want a large, elaborate SDXL pipeline built on a new canvas tab — aim for roughly 150 nodes total. Include: a base checkpoint pass (loader, positive/negative prompts, empty latent, KSampler, VAE decode), a refiner pass, a two-stage latent upscale with a second and third KSampler, four separate LoRA loaders chained on the model and clip, a ControlNet branch (ControlNet loader + apply + a LoadImage + preprocessor), several CLIPTextEncode variations combined with ConditioningCombine and ConditioningConcat nodes, multiple VAEDecode/VAEEncode round-trips, a face-detailer style crop-and-resample sub-graph, an image-composite branch, and a batch of SaveImage/PreviewImage nodes for each stage. Add plenty of Reroute nodes to keep the wiring tidy. Build it incrementally with add_node / set_node_widget and send_graph_mutation until the graph is large and complete." + ] + }, + { + "id": "t4", + "title": "Three successive edits + re-reads in one session", + "seedWorkflow": "edit-session-seed", + "prompts": [ + "The workflow on my canvas is a landscape generator with a LoRA. First edit: change the KSampler to 40 steps and switch the sampler to dpmpp_2m_sde. Then read back the graph and confirm the sampler settings.", + "Second edit: change the positive prompt to \"a dense bioluminescent alien jungle at night, volumetric fog, highly detailed\" and set the LoRA model strength to 1.0. Read the graph back and confirm both changes landed.", + "Third edit: add a second SaveImage node with the filename prefix \"jungle_alt\" wired from the same VAE decode output, then read the whole graph back one more time and give me a short summary of the final workflow." + ] + } + ] +} diff --git a/pyproject.toml b/pyproject.toml index 6cecbd4f..f91d096e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,9 @@ dependencies = [ "websocket-client", ] +# `bench` — the BE-2302 arm-B micro-edit benchmark runner (comfy_cli/bench). Only a +# live run needs the Anthropic SDK; the offline `--dry-run` + unit tests are stub-only. +optional-dependencies.bench = [ "anthropic>=0.40" ] optional-dependencies.dev = [ "jsonschema", "pre-commit", "pytest", "pytest-cov", "ruff" ] optional-dependencies.preview = [ "term-image>=0.7,<0.8" ] urls.Repository = "https://github.com/Comfy-Org/comfy-cli.git" @@ -74,6 +77,8 @@ include = [ "comfy_cli*" ] "comfy_cli.cql.data" = [ "*.yaml", "*.json" ] "comfy_cli.command.generate" = [ "spec/openapi.yml" ] +# Bench task matrix + offline object_info / seed-workflow fixtures. +"comfy_cli.bench" = [ "tasks.json", "fixtures/*.json" ] [tool.ruff] target-version = "py310" diff --git a/tests/comfy_cli/bench/__init__.py b/tests/comfy_cli/bench/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/comfy_cli/bench/test_run_arm_b.py b/tests/comfy_cli/bench/test_run_arm_b.py new file mode 100644 index 00000000..adc05fac --- /dev/null +++ b/tests/comfy_cli/bench/test_run_arm_b.py @@ -0,0 +1,282 @@ +"""Offline unit tests for the arm-B (CLI micro-edit proxy) bench runner. + +All tests use the committed object_info + seed fixtures and a stubbed model client — +NO live Anthropic calls, NO network — so CI stays offline. They cover the two pieces +BE-2309 calls out (tool dispatch + usage parsing) plus the end-to-end dry run and the +NDJSON row shape that must match arm A's ``arm-a.ndjson`` for report.mjs to consume it. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from comfy_cli.bench import run_arm_b as arm_b + +FIXTURES = Path(arm_b.__file__).resolve().parent / "fixtures" +OBJECT_INFO = FIXTURES / "object_info.json" +SEED = FIXTURES / "txt2img_seed.json" + +# The report.mjs / arm-a.ndjson shared row schema (extra keys like `proxy` are allowed). +REQUIRED_FIELDS = { + "arm", + "task", + "title", + "turn", + "msg_id", + "model", + "input_tokens", + "output_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", + "cost_usd", + "sdk_cost_usd", + "num_turns", + "duration_ms", + "duration_api_ms", + "wall_ms", + "tool_call_count", + "tool_input_bytes_total", + "tool_input_bytes_max", + "tool_output_bytes_total", + "tool_output_bytes_max", + "tool_payload_bytes_total", + "tool_payload_bytes_max", + "compactions", +} + + +# --------------------------------------------------------------------------- # +# usage parsing +# --------------------------------------------------------------------------- # + + +def test_parse_usage_from_dict(): + usage = { + "input_tokens": 100, + "output_tokens": 50, + "cache_read_input_tokens": 2000, + "cache_creation_input_tokens": 300, + } + assert arm_b.parse_usage(usage) == usage + + +def test_parse_usage_from_object(): + class Usage: + input_tokens = 12 + output_tokens = 7 + cache_read_input_tokens = 900 + cache_creation_input_tokens = 40 + + assert arm_b.parse_usage(Usage()) == { + "input_tokens": 12, + "output_tokens": 7, + "cache_read_input_tokens": 900, + "cache_creation_input_tokens": 40, + } + + +def test_parse_usage_missing_and_none_default_to_zero(): + # Partial block (no cache fields) + None → all missing counts are 0, never a crash. + assert arm_b.parse_usage({"input_tokens": 5}) == { + "input_tokens": 5, + "output_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + } + assert arm_b.parse_usage(None) == dict.fromkeys(arm_b.TOKEN_FIELDS, 0) + + +def test_parse_usage_ignores_bool_and_nonnumeric(): + assert arm_b.parse_usage({"input_tokens": True, "output_tokens": "x"}) == dict.fromkeys(arm_b.TOKEN_FIELDS, 0) + + +# --------------------------------------------------------------------------- # +# pricing — must match comfy-inapp-agent/agent-server/usage.mjs exactly +# --------------------------------------------------------------------------- # + + +def test_estimate_cost_matches_usage_mjs_pricing(): + # Opus 4.8: $5/M in, $25/M out, cache read 0.1x, cache write 1.25x. + usage = { + "input_tokens": 3950, + "output_tokens": 610, + "cache_read_input_tokens": 18200, + "cache_creation_input_tokens": 3000, + } + expected = (3950 * 5 + 610 * 25 + 18200 * 5 * 0.1 + 3000 * 5 * 1.25) / 1e6 + assert arm_b.estimate_cost_usd("claude-opus-4-8", usage) == pytest.approx(expected) + + +def test_estimate_cost_unknown_model_is_none(): + assert arm_b.estimate_cost_usd("gpt-4o", {"input_tokens": 1000}) is None + assert arm_b.estimate_cost_usd(None, {}) is None + + +def test_resolve_pricing_longest_prefix(): + # A dated snapshot resolves to its family entry by longest prefix. + entry = arm_b.resolve_pricing("claude-opus-4-8-20260101") + assert entry is not None and entry[0] == "claude-opus-4-8" + + +# --------------------------------------------------------------------------- # +# tool dispatch — real CQL engine against the committed fixtures +# --------------------------------------------------------------------------- # + + +@pytest.fixture +def dispatcher(tmp_path): + wf = tmp_path / "wf.json" + wf.write_text(SEED.read_text(encoding="utf-8"), encoding="utf-8") + return arm_b.ToolDispatcher(OBJECT_INFO, wf, tmp_path / "variants") + + +def test_dispatch_slots_lists_addresses(dispatcher): + result, is_error = dispatcher.dispatch("slots", {}) + assert not is_error + addrs = {s["address"] for s in result["slots"]} + assert "6.text" in addrs and "3.steps" in addrs + assert result["count"] == len(result["slots"]) + + +def test_dispatch_cql_found_and_not_found(dispatcher): + ok, err = dispatcher.dispatch("cql", {"node_type": "KSampler"}) + assert not err and ok["found"] is True and ok["schema"]["id"] == "KSampler" + + miss, err = dispatcher.dispatch("cql", {"node_type": "LoraLoader"}) + assert not err and miss["found"] is False and "suggestions" in miss + + +def test_dispatch_cql_no_arg_lists_catalog(dispatcher): + result, err = dispatcher.dispatch("cql", {}) + assert not err and "CheckpointLoaderSimple" in result["available_node_types"] + + +def test_dispatch_set_slot_mutates_file_on_disk(dispatcher): + result, is_error = dispatcher.dispatch("set_slot", {"overrides": {"6.text": "a brand new prompt"}}) + assert not is_error and result["applied"] == ["6.text"] + # The write is a genuine on-disk round-trip: a fresh `slots` read sees the change. + slots, _ = dispatcher.dispatch("slots", {}) + positive = next(s for s in slots["slots"] if s["address"] == "6.text") + assert positive["current_value"] == "a brand new prompt" + + +def test_dispatch_set_slot_bad_address_is_soft_error(dispatcher): + result, is_error = dispatcher.dispatch("set_slot", {"overrides": {"999.nope": 1}}) + assert is_error and "error" in result # surfaced to the model, never raised + + +def test_dispatch_set_slot_requires_overrides(dispatcher): + result, is_error = dispatcher.dispatch("set_slot", {}) + assert is_error and "error" in result + + +def test_dispatch_vary_produces_variants(dispatcher): + result, is_error = dispatcher.dispatch("vary", {"slots": {"3.seed": [1, 2, 3]}}) + assert not is_error and result["count"] == 3 and len(result["written"]) == 3 + + +def test_dispatch_vary_mismatched_lengths_is_error(dispatcher): + result, is_error = dispatcher.dispatch("vary", {"slots": {"3.seed": [1, 2], "3.steps": [10]}}) + assert is_error and "error" in result + + +def test_dispatch_unknown_tool_is_error(dispatcher): + result, is_error = dispatcher.dispatch("nonesuch", {}) + assert is_error and "unknown tool" in result["error"] + + +# --------------------------------------------------------------------------- # +# driver loop with the stub client +# --------------------------------------------------------------------------- # + + +def test_stub_driver_runs_a_turn_and_accumulates_telemetry(tmp_path): + wf = tmp_path / "wf.json" + wf.write_text(SEED.read_text(encoding="utf-8"), encoding="utf-8") + disp = arm_b.ToolDispatcher(OBJECT_INFO, wf, tmp_path / "v") + driver = arm_b._StubDriver(arm_b.StubClient(), model=arm_b.DEFAULT_MODEL) + + tel = driver.run_turn_for(disp, "t2", 1, "change the prompt") + # t2 script: slots + set_slot + final text = 3 API calls, 2 tool calls. + assert tel.num_api_calls == 3 + assert tel.tool_call_count == 2 + assert tel.input_tokens > 0 and tel.tool_input_bytes and tel.tool_output_bytes + # The set_slot in the script really edited the temp workflow. + assert json.loads(wf.read_text())["nodes"] # still valid frontend JSON + + +def test_build_row_has_full_schema_and_proxy_label(): + tel = arm_b.TurnTelemetry() + tel.add_usage( + {"input_tokens": 100, "output_tokens": 20, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0} + ) + tel.add_tool(50, 200) + task = {"id": "t2", "title": "Micro-edit one parameter in a seed workflow"} + row = arm_b.build_row(task=task, turn=1, model=arm_b.DEFAULT_MODEL, tel=tel) + assert REQUIRED_FIELDS <= set(row.keys()) + assert row["arm"] == "B" and row["proxy"] is True + assert row["tool_payload_bytes_total"] == 250 and row["tool_payload_bytes_max"] == 250 + assert row["cost_usd"] == row["sdk_cost_usd"] + + +# --------------------------------------------------------------------------- # +# end-to-end dry run — the acceptance criterion +# --------------------------------------------------------------------------- # + + +def test_dry_run_writes_well_formed_ndjson(tmp_path): + out = tmp_path / "arm-b.ndjson" + rows = arm_b.run( + dry_run=True, + out_path=out, + object_info_path=OBJECT_INFO, + tasks_path=arm_b.TASKS_PATH, + model=arm_b.DEFAULT_MODEL, + work_dir=tmp_path / "work", + ) + # One row per (task, turn): t1, t2, t3, t4x3 = 6. + assert len(rows) == 6 + file_rows = [json.loads(line) for line in out.read_text(encoding="utf-8").splitlines() if line.strip()] + assert file_rows == rows + for r in file_rows: + assert REQUIRED_FIELDS <= set(r.keys()) + assert r["arm"] == "B" and r["proxy"] is True + assert isinstance(r["cost_usd"], int | float) and r["cost_usd"] >= 0 + + by_task = {(r["task"], r["turn"]): r for r in file_rows} + # t3 is the documented ceiling, recorded as a RESULT. + assert by_task[("t3", 1)]["outcome"] == "ceiling" + assert by_task[("t3", 1)]["note"] + # t4 is a 3-turn session; cache-read tokens grow turn-over-turn (transcript regrowth). + assert ( + by_task[("t4", 1)]["cache_read_input_tokens"] + < by_task[("t4", 2)]["cache_read_input_tokens"] + < by_task[("t4", 3)]["cache_read_input_tokens"] + ) + + +def test_dry_run_only_filter(tmp_path): + out = tmp_path / "arm-b.ndjson" + rows = arm_b.run( + dry_run=True, + out_path=out, + object_info_path=OBJECT_INFO, + tasks_path=arm_b.TASKS_PATH, + model=arm_b.DEFAULT_MODEL, + work_dir=tmp_path / "work", + only=["t2"], + ) + assert [r["task"] for r in rows] == ["t2"] + + +def test_tasks_json_prompts_are_present_and_nonempty(): + tasks = arm_b.load_tasks() + ids = [t["id"] for t in tasks] + assert ids == ["t1", "t2", "t3", "t4"] + assert len(next(t for t in tasks if t["id"] == "t4")["prompts"]) == 3 + for t in tasks: + for p in t["prompts"]: + assert isinstance(p, str) and p.strip() From e05522e49770e0b17621aca73fc02c360d191eaa Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Thu, 2 Jul 2026 19:14:01 -0700 Subject: [PATCH 2/2] fix(bench): harden arm-B runner per Cursor review (BE-2309) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address the Cursor consolidated-panel findings on PR #490: - fixtures: edit_session_seed KSampler was missing the control_after_generate widget value, shifting every slot after `seed` off by one (steps/cfg/sampler read + written to the wrong indices on t4). Add the 7th value; add a regression test asserting the slots align. - run_turn: always answer emitted tool_use blocks before breaking, and stop the task's remaining turns when a turn is truncated (non-terminal stop reason or MAX_API_CALLS_PER_TURN exhaustion) — keeps the shared t4 transcript valid and stamps such turns outcome="truncated" instead of a false "ok". - fail-fast: unknown --model (no pricing entry) now errors instead of silently pricing every turn at $0; unrecognized non-null seedWorkflow errors instead of silently running against the wrong seed graph. - safety: validate task ids before interpolating them into filesystem paths (reject separators/`..`); cap `vary` variant count at MAX_VARIANTS and namespace each vary call's output dir so a second call can't overwrite the first; return a generic tool-error message (log detail locally) so raw exception strings / local paths aren't sent to the Anthropic API. - telemetry: count tool payload bytes as UTF-8 (ensure_ascii=False) instead of escaped-ASCII lengths so the byte comparison against arm A isn't skewed. Co-Authored-By: Claude Opus 4.8 --- .../bench/fixtures/edit_session_seed.json | 1 + comfy_cli/bench/run_arm_b.py | 147 +++++++++++++++--- tests/comfy_cli/bench/test_run_arm_b.py | 96 ++++++++++++ 3 files changed, 219 insertions(+), 25 deletions(-) diff --git a/comfy_cli/bench/fixtures/edit_session_seed.json b/comfy_cli/bench/fixtures/edit_session_seed.json index 603184fc..94b07592 100644 --- a/comfy_cli/bench/fixtures/edit_session_seed.json +++ b/comfy_cli/bench/fixtures/edit_session_seed.json @@ -115,6 +115,7 @@ }, "widgets_values": [ 123456789, + "fixed", 20, 7.0, "euler", diff --git a/comfy_cli/bench/run_arm_b.py b/comfy_cli/bench/run_arm_b.py index 8dd232e4..ce817f17 100644 --- a/comfy_cli/bench/run_arm_b.py +++ b/comfy_cli/bench/run_arm_b.py @@ -39,6 +39,7 @@ import argparse import json +import logging import shutil import sys import time @@ -47,6 +48,8 @@ from pathlib import Path from typing import Any +logger = logging.getLogger(__name__) + # --------------------------------------------------------------------------- # Pricing — mirror of comfy-inapp-agent/agent-server/usage.mjs PRICING so the two # arms price a turn identically. Values are USD per million tokens. @@ -66,6 +69,10 @@ DEFAULT_MODEL = "claude-opus-4-8" +# Upper bound on variants a single `vary` tool call may produce — guards against an +# unbounded model-supplied list length exhausting memory/disk. +MAX_VARIANTS = 64 + # Token fields the runner captures from the Anthropic ``usage`` block, emitted # verbatim so report.mjs can recompute cost from them. TOKEN_FIELDS = ( @@ -219,6 +226,7 @@ def __init__(self, object_info_path: str | Path, workflow_path: str | Path, vari self.graph = Graph.load(input_path=str(object_info_path)) self.workflow_path = Path(workflow_path) self.variant_dir = Path(variant_dir) + self._vary_calls = 0 # namespaces each vary invocation's output files def _load_workflow(self) -> dict[str, Any]: return json.loads(self.workflow_path.read_text(encoding="utf-8")) @@ -235,8 +243,12 @@ def dispatch(self, name: str, tool_input: dict[str, Any]) -> tuple[dict[str, Any if handler is None: return {"error": f"unknown tool {name!r}"}, True return handler(tool_input or {}) - except Exception as e: # dispatch must be total — surface the error to the model - return {"error": f"{type(e).__name__}: {e}"}, True + except Exception: # dispatch must be total — never raise into the loop + # Log full detail locally, but return only a generic message: the result + # is sent to the external Anthropic API on a live run and the raw exception + # string can embed absolute filesystem paths / other local state. + logger.exception("tool %r raised while dispatching", name) + return {"error": f"internal error while running tool {name!r}"}, True # -- read tools -- @@ -288,18 +300,30 @@ def _vary(self, inp: dict[str, Any]) -> tuple[dict[str, Any], bool]: if len(set(lengths.values())) != 1: return {"error": f"all `slots` lists must be the same length; got {lengths}."}, True n = next(iter(lengths.values())) + if n < 1: + return {"error": "vary requires at least one value per slot."}, True + if n > MAX_VARIANTS: + # A model-supplied list length is unbounded; deep-copying the workflow n + # times would exhaust memory/disk/inodes. Cap it so a runaway tool call fails + # loudly instead of hanging. + return {"error": f"vary is capped at {MAX_VARIANTS} variants per call; requested {n}."}, True variations = [{a: v[i] for a, v in by_addr.items()} for i in range(n)] wf = self._load_workflow() try: workflows, warnings = self.graph.expand_variations(wf, variations) except ValueError as e: return {"error": str(e)}, True - self.variant_dir.mkdir(parents=True, exist_ok=True) + # Namespace each vary call's files so a second call in the same task can't + # silently overwrite the first call's variants (both would restart at _000). + call_idx = self._vary_calls + self._vary_calls += 1 + out_dir = self.variant_dir / f"call_{call_idx:03d}" + out_dir.mkdir(parents=True, exist_ok=True) written = [] for i, w in enumerate(workflows): - target = self.variant_dir / f"{self.workflow_path.stem}_{i:03d}.json" + target = out_dir / f"{self.workflow_path.stem}_{i:03d}.json" target.write_text(json.dumps(w, indent=2), encoding="utf-8") - written.append(target.name) + written.append(f"{out_dir.name}/{target.name}") return {"count": len(workflows), "written": written, "warnings": warnings}, False @@ -328,6 +352,9 @@ class TurnTelemetry: tool_input_bytes: list[int] = field(default_factory=list) tool_output_bytes: list[int] = field(default_factory=list) duration_ms: float = 0.0 + # Turn completion status: "ok" when the model ended cleanly, "truncated" when the + # turn was cut short (non-terminal stop reason or MAX_API_CALLS_PER_TURN exhaustion). + stop_status: str = "ok" def add_usage(self, usage: dict[str, int]) -> None: for f in TOKEN_FIELDS: @@ -434,6 +461,7 @@ def run_turn( messages = [] tel = TurnTelemetry() messages.append({"role": "user", "content": prompt}) + completed = False for _ in range(MAX_API_CALLS_PER_TURN): t0 = time.perf_counter() resp = self._create(messages) @@ -446,24 +474,45 @@ def run_turn( content = _battr(resp, "content") or [] messages.append({"role": "assistant", "content": content}) tool_uses = [b for b in content if _battr(b, "type") == "tool_use"] - if _battr(resp, "stop_reason") != "tool_use" or not tool_uses: + stop_reason = _battr(resp, "stop_reason") + + # Always answer every tool_use the model emitted — even when it also signalled + # a terminal stop_reason. Leaving a tool_use unanswered makes the transcript + # invalid, and since a task's turns (t4) share one `messages` list, the next + # turn's create() would 400 on the dangling tool_use. + if tool_uses: + tool_results = [] + for tu in tool_uses: + tu_input = _battr(tu, "input") or {} + result, is_error = dispatcher.dispatch(_battr(tu, "name"), tu_input) + result_json = json.dumps(result, ensure_ascii=False) + # Measure true UTF-8 payload bytes (not escaped-ASCII lengths) so the + # byte comparison against arm A isn't skewed by unicode-heavy payloads. + tel.add_tool( + len(json.dumps(tu_input, ensure_ascii=False).encode("utf-8")), + len(result_json.encode("utf-8")), + ) + tool_results.append( + { + "type": "tool_result", + "tool_use_id": _battr(tu, "id"), + "content": result_json, + "is_error": is_error, + } + ) + messages.append({"role": "user", "content": tool_results}) + + if stop_reason != "tool_use" or not tool_uses: + # The model stopped. "end_turn"/"stop_sequence" are clean completions; + # anything else (e.g. "max_tokens") means the turn was cut short. + completed = True + if stop_reason not in ("end_turn", "stop_sequence", "tool_use"): + tel.stop_status = "truncated" break - - tool_results = [] - for tu in tool_uses: - tu_input = _battr(tu, "input") or {} - result, is_error = dispatcher.dispatch(_battr(tu, "name"), tu_input) - result_json = json.dumps(result) - tel.add_tool(len(json.dumps(tu_input)), len(result_json)) - tool_results.append( - { - "type": "tool_result", - "tool_use_id": _battr(tu, "id"), - "content": result_json, - "is_error": is_error, - } - ) - messages.append({"role": "user", "content": tool_results}) + if not completed: + # Exhausted MAX_API_CALLS_PER_TURN still mid-tool-loop: the transcript ends on a + # dangling tool_result (user) turn, so the task must not continue onto it. + tel.stop_status = "truncated" return tel @@ -477,7 +526,16 @@ def _seed_for_task(task: dict[str, Any], fixtures_dir: Path) -> Path: """ name = task.get("seedWorkflow") mapping = {"txt2img-seed": "txt2img_seed.json", "edit-session-seed": "edit_session_seed.json"} - return fixtures_dir / mapping.get(name, "txt2img_seed.json") + if name is None: + # t1/t3 build from scratch; arm B seeds from the txt2img template and micro-edits it. + return fixtures_dir / "txt2img_seed.json" + if name not in mapping: + # An unrecognized non-null seed id (typo / new upstream id) must fail loudly rather + # than silently run against the wrong seed graph. + raise ValueError( + f"unknown seedWorkflow {name!r} for task {task.get('id')!r}; expected one of {sorted(mapping)}" + ) + return fixtures_dir / mapping[name] # t3 is the documented ceiling: arm B's slots/set-slot/vary can only EDIT an existing @@ -497,6 +555,22 @@ def _task_outcome(task_id: str) -> tuple[str, str | None]: return "ok", None +def _resolve_outcome(task_id: str, tel: TurnTelemetry) -> tuple[str, str | None]: + """Combine the per-task stamp with the turn's actual completion status. + + A truncated turn (non-terminal stop reason or MAX_API_CALLS_PER_TURN exhaustion) + must NOT be emitted as a successful row — that would corrupt the arm-A/arm-B + comparison this runner exists to produce. + """ + if tel.stop_status == "truncated": + return ( + "truncated", + "turn cut short (non-terminal stop reason or MAX_API_CALLS_PER_TURN exhausted); " + "remaining turns skipped to keep the shared transcript valid", + ) + return _task_outcome(task_id) + + def run_task( task: dict[str, Any], *, @@ -519,8 +593,10 @@ def run_task( messages: list[dict[str, Any]] = [] for i, prompt in enumerate(task["prompts"], start=1): tel = driver.run_turn(dispatcher, prompt, messages) - outcome, note = _task_outcome(task["id"]) + outcome, note = _resolve_outcome(task["id"], tel) on_row(build_row(task=task, turn=i, model=driver.model, tel=tel, outcome=outcome, note=note)) + if tel.stop_status == "truncated": + break # shared transcript is now invalid; don't run further turns on it # --------------------------------------------------------------------------- @@ -759,8 +835,10 @@ def run_task_stub( messages: list[dict[str, Any]] = [] for i, prompt in enumerate(task["prompts"], start=1): tel = driver.run_turn_for(dispatcher, task["id"], i, prompt, messages) - outcome, note = _task_outcome(task["id"]) + outcome, note = _resolve_outcome(task["id"], tel) on_row(build_row(task=task, turn=i, model=driver.model, tel=tel, outcome=outcome, note=note)) + if tel.stop_status == "truncated": + break # shared transcript is now invalid; don't run further turns on it # --------------------------------------------------------------------------- @@ -797,11 +875,25 @@ def build_live_client() -> Any: DEFAULT_OUT = _HERE / "out" / "arm-b.ndjson" +def _validate_task_id(task_id: Any) -> str: + """Reject task ids unsafe to interpolate into a filesystem path. + + ``task['id']`` flows straight into scratch-file/variant-dir names; a crafted id with + a path separator or ``..`` would let a malicious ``--tasks`` file escape ``work_dir``. + """ + if not isinstance(task_id, str) or not task_id: + raise ValueError(f"task id must be a non-empty string, got {task_id!r}") + if task_id in (".", "..") or "/" in task_id or "\\" in task_id or "\x00" in task_id: + raise ValueError(f"task id {task_id!r} may not contain path separators or '..'") + return task_id + + def load_tasks(path: str | Path = TASKS_PATH) -> list[dict[str, Any]]: data = json.loads(Path(path).read_text(encoding="utf-8")) tasks = data["tasks"] # Normalize: every task has a `prompts` list (matching child 1's taskPrompts()). for t in tasks: + _validate_task_id(t.get("id")) if "prompts" not in t or not isinstance(t["prompts"], list): t["prompts"] = [t["prompt"]] if t.get("prompt") else [] return tasks @@ -823,6 +915,11 @@ def run( only: list[str] | None = None, ) -> list[dict[str, Any]]: """Run the whole matrix, write NDJSON to ``out_path``, and return the rows.""" + if resolve_pricing(model) is None: + # Fail fast: an unknown model would price every turn at $0 (estimate_cost_usd → None + # → 0.0), making the run look free and silently skewing the dollar comparison. + known = ", ".join(entry[0] for entry in PRICING) + raise ValueError(f"model {model!r} has no pricing entry; add one to PRICING (known prefixes: {known}).") tasks = load_tasks(tasks_path) if only: tasks = [t for t in tasks if t["id"] in set(only)] diff --git a/tests/comfy_cli/bench/test_run_arm_b.py b/tests/comfy_cli/bench/test_run_arm_b.py index adc05fac..b5afa874 100644 --- a/tests/comfy_cli/bench/test_run_arm_b.py +++ b/tests/comfy_cli/bench/test_run_arm_b.py @@ -188,6 +188,49 @@ def test_dispatch_unknown_tool_is_error(dispatcher): assert is_error and "unknown tool" in result["error"] +def test_dispatch_vary_over_cap_is_error(dispatcher): + # An unbounded model-supplied list length is capped so a runaway call fails loudly. + result, is_error = dispatcher.dispatch("vary", {"slots": {"3.seed": list(range(arm_b.MAX_VARIANTS + 1))}}) + assert is_error and "capped" in result["error"] + + +def test_dispatch_vary_calls_do_not_overwrite(dispatcher): + # Two vary calls in one task must not clobber each other's variant files. + r1, e1 = dispatcher.dispatch("vary", {"slots": {"3.seed": [1, 2]}}) + r2, e2 = dispatcher.dispatch("vary", {"slots": {"3.seed": [3, 4]}}) + assert not e1 and not e2 + assert set(r1["written"]).isdisjoint(r2["written"]) + # Both calls' files coexist on disk (namespaced per call). + written = set(r1["written"]) | set(r2["written"]) + assert len(list((dispatcher.variant_dir).rglob("*.json"))) == len(written) == 4 + + +def test_dispatch_catch_all_error_is_generic(dispatcher, monkeypatch): + # An unexpected exception must not leak its raw message (possibly a local path) to the model. + def boom(*_a, **_k): + raise RuntimeError("/secret/local/path leaked") + + monkeypatch.setattr(dispatcher, "_slots", boom) + result, is_error = dispatcher.dispatch("slots", {}) + assert is_error and "secret" not in result["error"] and "internal error" in result["error"] + + +def test_edit_session_seed_ksampler_slots_align(): + # Regression: the KSampler widgets_values must include control_after_generate so slots + # map positionally (steps=20, cfg=7.0), not off-by-one. + import tempfile + + with tempfile.TemporaryDirectory() as td: + wf = Path(td) / "wf.json" + wf.write_text((FIXTURES / "edit_session_seed.json").read_text(encoding="utf-8"), encoding="utf-8") + disp = arm_b.ToolDispatcher(OBJECT_INFO, wf, Path(td) / "v") + slots, _ = disp.dispatch("slots", {}) + by_addr = {s["address"]: s["current_value"] for s in slots["slots"]} + assert by_addr["3.steps"] == 20 + assert by_addr["3.cfg"] == 7.0 + assert by_addr["3.sampler_name"] == "euler" + + # --------------------------------------------------------------------------- # # driver loop with the stub client # --------------------------------------------------------------------------- # @@ -280,3 +323,56 @@ def test_tasks_json_prompts_are_present_and_nonempty(): for t in tasks: for p in t["prompts"]: assert isinstance(p, str) and p.strip() + + +# --------------------------------------------------------------------------- # +# input validation + fail-fast guards +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("bad", ["../evil", "a/b", "a\\b", "..", ".", "", None, 5]) +def test_validate_task_id_rejects_unsafe(bad): + with pytest.raises(ValueError): + arm_b._validate_task_id(bad) + + +def test_validate_task_id_accepts_normal(): + assert arm_b._validate_task_id("t4") == "t4" + + +def test_seed_for_task_unknown_seed_fails_fast(): + with pytest.raises(ValueError): + arm_b._seed_for_task({"id": "tx", "seedWorkflow": "bogus-seed"}, FIXTURES) + # None seed (t1/t3 build-from-scratch) still resolves to the txt2img template. + assert arm_b._seed_for_task({"id": "t1", "seedWorkflow": None}, FIXTURES).name == "txt2img_seed.json" + + +def test_run_unknown_model_fails_fast(tmp_path): + with pytest.raises(ValueError): + arm_b.run( + dry_run=True, + out_path=tmp_path / "out.ndjson", + object_info_path=OBJECT_INFO, + tasks_path=arm_b.TASKS_PATH, + model="gpt-4o", + work_dir=tmp_path / "work", + ) + + +def test_truncated_turn_marks_outcome_and_stops(tmp_path): + # A turn cut short (non-terminal stop reason) is stamped "truncated" and the task's + # remaining turns are skipped so the shared transcript never goes invalid. + plan = {("t4", 1): [([arm_b._text("cut off mid-thought")], "max_tokens", arm_b._usage(100, 20, 0, 0))]} + driver = arm_b._StubDriver(arm_b.StubClient(plan), model=arm_b.DEFAULT_MODEL) + task = {"id": "t4", "title": "edit session", "seedWorkflow": "edit-session-seed", "prompts": ["a", "b", "c"]} + rows: list[dict] = [] + arm_b.run_task_stub( + task, + driver=driver, + fixtures_dir=FIXTURES, + object_info_path=OBJECT_INFO, + work_dir=tmp_path, + on_row=rows.append, + ) + assert len(rows) == 1 # turns 2 and 3 skipped + assert rows[0]["outcome"] == "truncated" and rows[0]["note"]