|
| 1 | +"""TDPilot API — ``tool_batch`` handler (Phase 2.1). |
| 2 | +
|
| 3 | +Run multiple TDPilot tool calls in one round trip. The win is |
| 4 | +LLM-round-trip latency: instead of issuing N tool_use blocks and |
| 5 | +paying N model→server→model cycles, the agent issues one |
| 6 | +``tool_batch`` and gets all N results back together. |
| 7 | +
|
| 8 | +Sub-calls execute SEQUENTIALLY on the cook thread because TD's |
| 9 | +Python API isn't thread-safe — ``ThreadPoolExecutor.submit`` |
| 10 | +doesn't help here, the per-tool latency is unchanged. What we save |
| 11 | +is the chain of ``tool_use → tool_result → next-think → next |
| 12 | +tool_use`` that the model would otherwise traverse. |
| 13 | +
|
| 14 | +Failure handling: a sub-call that returns ``{"error": ...}`` does |
| 15 | +NOT abort the batch. Each result is reported in |
| 16 | +``results[i].error`` and the rest still run. This matches the way |
| 17 | +``recipe_replay`` was always *supposed* to behave when the agent |
| 18 | +runs heterogeneous read-only lookups. |
| 19 | +
|
| 20 | +Hard caps: |
| 21 | + max 8 sub-calls (matches the schema's maxItems). |
| 22 | + nested ``tool_batch`` is rejected per sub-call (cheap recursion |
| 23 | + guard so a confused agent can't fork-bomb the dispatcher). |
| 24 | +""" |
| 25 | + |
| 26 | +from __future__ import annotations |
| 27 | + |
| 28 | +import time |
| 29 | +from typing import Any |
| 30 | + |
| 31 | +MAX_BATCH_SIZE = 8 |
| 32 | + |
| 33 | + |
| 34 | +def _resolve_raw_dispatcher() -> Any: |
| 35 | + """Walk COMP → extension → runtime to find the cook-thread-bypass |
| 36 | + dispatcher (same pattern as ``handle_recipe_replay``). |
| 37 | +
|
| 38 | + Returns ``None`` if any step fails — caller surfaces a clear error. |
| 39 | + """ |
| 40 | + try: |
| 41 | + comp = parent() # type: ignore[name-defined] |
| 42 | + except NameError: |
| 43 | + return None |
| 44 | + if comp is None: |
| 45 | + return None |
| 46 | + ext_dat = comp.op("tdpilot_api_extension") |
| 47 | + if ext_dat is None: |
| 48 | + return None |
| 49 | + try: |
| 50 | + ext = ext_dat.module.get_extension(comp) |
| 51 | + return ext._runtime._raw_dispatcher |
| 52 | + except Exception: |
| 53 | + return None |
| 54 | + |
| 55 | + |
| 56 | +def handle_tool_batch(body: dict) -> dict: |
| 57 | + """Dispatch a list of tool calls and return all results. |
| 58 | +
|
| 59 | + Body schema: |
| 60 | + {"calls": [{"tool": str, "args": dict}, ...]} |
| 61 | +
|
| 62 | + Returns: |
| 63 | + {"ok": True, |
| 64 | + "count": int, |
| 65 | + "results": [ |
| 66 | + {"tool": str, |
| 67 | + "ok": bool, |
| 68 | + "result": dict | None, |
| 69 | + "error": str | None, |
| 70 | + "elapsed_ms": int}, |
| 71 | + ... |
| 72 | + ]} |
| 73 | +
|
| 74 | + On invalid body shape returns ``{"error": str}`` instead. |
| 75 | + """ |
| 76 | + if not isinstance(body, dict): |
| 77 | + return {"error": "tool_batch body must be a dict"} |
| 78 | + |
| 79 | + calls = body.get("calls") |
| 80 | + if not isinstance(calls, list) or not calls: |
| 81 | + return {"error": "tool_batch requires non-empty 'calls' list"} |
| 82 | + |
| 83 | + if len(calls) > MAX_BATCH_SIZE: |
| 84 | + return {"error": (f"tool_batch capped at {MAX_BATCH_SIZE} sub-calls (received {len(calls)})")} |
| 85 | + |
| 86 | + raw_dispatcher = _resolve_raw_dispatcher() |
| 87 | + if raw_dispatcher is None: |
| 88 | + return {"error": "tool_batch could not access the runtime dispatcher"} |
| 89 | + |
| 90 | + results: list[dict] = [] |
| 91 | + for i, call in enumerate(calls): |
| 92 | + if not isinstance(call, dict): |
| 93 | + results.append( |
| 94 | + { |
| 95 | + "tool": None, |
| 96 | + "ok": False, |
| 97 | + "result": None, |
| 98 | + "error": f"call[{i}] is not an object", |
| 99 | + "elapsed_ms": 0, |
| 100 | + } |
| 101 | + ) |
| 102 | + continue |
| 103 | + tool_name = call.get("tool") |
| 104 | + if not isinstance(tool_name, str) or not tool_name.strip(): |
| 105 | + results.append( |
| 106 | + { |
| 107 | + "tool": tool_name, |
| 108 | + "ok": False, |
| 109 | + "result": None, |
| 110 | + "error": f"call[{i}].tool is missing or not a string", |
| 111 | + "elapsed_ms": 0, |
| 112 | + } |
| 113 | + ) |
| 114 | + continue |
| 115 | + if tool_name == "tool_batch": |
| 116 | + results.append( |
| 117 | + { |
| 118 | + "tool": tool_name, |
| 119 | + "ok": False, |
| 120 | + "result": None, |
| 121 | + "error": "Nested tool_batch is not allowed", |
| 122 | + "elapsed_ms": 0, |
| 123 | + } |
| 124 | + ) |
| 125 | + continue |
| 126 | + |
| 127 | + tool_args = call.get("args") or {} |
| 128 | + if not isinstance(tool_args, dict): |
| 129 | + results.append( |
| 130 | + { |
| 131 | + "tool": tool_name, |
| 132 | + "ok": False, |
| 133 | + "result": None, |
| 134 | + "error": f"call[{i}].args must be an object", |
| 135 | + "elapsed_ms": 0, |
| 136 | + } |
| 137 | + ) |
| 138 | + continue |
| 139 | + |
| 140 | + t_start = time.monotonic() |
| 141 | + try: |
| 142 | + result = raw_dispatcher(tool_name, tool_args) |
| 143 | + except Exception as exc: # noqa: BLE001 — surface as per-call error |
| 144 | + elapsed_ms = int((time.monotonic() - t_start) * 1000) |
| 145 | + results.append( |
| 146 | + { |
| 147 | + "tool": tool_name, |
| 148 | + "ok": False, |
| 149 | + "result": None, |
| 150 | + "error": f"{type(exc).__name__}: {exc}", |
| 151 | + "elapsed_ms": elapsed_ms, |
| 152 | + } |
| 153 | + ) |
| 154 | + continue |
| 155 | + elapsed_ms = int((time.monotonic() - t_start) * 1000) |
| 156 | + |
| 157 | + is_error = isinstance(result, dict) and "error" in result |
| 158 | + results.append( |
| 159 | + { |
| 160 | + "tool": tool_name, |
| 161 | + "ok": not is_error, |
| 162 | + "result": result if not is_error else None, |
| 163 | + "error": (result.get("error") if is_error else None), |
| 164 | + "elapsed_ms": elapsed_ms, |
| 165 | + } |
| 166 | + ) |
| 167 | + |
| 168 | + return { |
| 169 | + "ok": True, |
| 170 | + "count": len(results), |
| 171 | + "results": results, |
| 172 | + } |
0 commit comments