diff --git a/packages/opencode-agents/README.md b/packages/opencode-agents/README.md index 9037e5d5..869bf9a1 100644 --- a/packages/opencode-agents/README.md +++ b/packages/opencode-agents/README.md @@ -36,4 +36,26 @@ asyncio.run(main()) `cwd` is the working directory, `config` is the OpenCode config file, and `env` is the child environment. The handler sets `OPENCODE_CONFIG` from `config`, replaces any `OPENCODE_SERVER_USERNAME` in `env` with `opencode`, replaces any `OPENCODE_SERVER_PASSWORD` with a generated secret, and binds `opencode serve` to loopback. Requests to that server use that username and secret as basic auth. A caller does not pass a URL, an HTTP client, or a transport. `command` overrides the `opencode` executable; a test passes a fake executable there. -Entering the handler starts `opencode serve` and returns once it is healthy. The process stops when the context exits, after `config().stream()` returns. Put a wall-clock limit around `config().stream()`, not around the context. A judge routed here runs as a full OpenCode agent: the server's tools, in `cwd`, apply to the judge prompt. There is no tool-less judge mode. The handler yields AG-UI text and tool events. `config().stream()` yields `chunk` then `done`. +Entering the handler starts `opencode serve` and returns once it is healthy. The process stops when the context exits, after `config().stream()` returns. Put a wall-clock limit around `config().stream()`, not around the context. The handler yields AG-UI text and tool events. `config().stream()` yields `chunk` then `done`. + +## Tools + +The AI Config's `tools`, filtered to the keys also present in `tool_handlers`, is the only catalog the model sees for one call — every other tool on the running `opencode serve`, including every built-in, is denied. A variation with no tools (and every judge, which always calls with `tool_handlers=None`) gets a model with no tools at all. + +```python +from launchdarkly_ai_opencode_agents import OpenCodeBash, OpenCodeRead + +response = await config( + key="my-agent", + handler=handler, + tool_handlers={ + "run-shell": OpenCodeBash, # an OpenCode built-in, allowed by name + "read-file": OpenCodeRead, + "get-weather": get_weather, # a plain callable, served from a per-call bridge + }, +).invoke(user_input, context) +``` + +A `NativeTool` (see `OpenCodeBash`, `OpenCodeRead`, and the rest of this package's sentinels) names an OpenCode built-in or a server-configured MCP tool directly — OpenCode already has it, so the handler only allows it for this call's session. A plain callable has no route into OpenCode otherwise: `POST /session/{id}/message`'s own `tools` field is a deprecated boolean permission map, not a place to describe one. Each call with at least one callable starts its own loopback MCP server, registers it on the running `opencode serve` with `POST /mcp`, and disconnects it when the call ends. This is unit-tested directly (`test_tool_bridge.py`) and through the handler against a fake `opencode serve` (`test_tool_contract.py`); it has not been exercised against a real OpenCode MCP client. + +`chat` and `execute_tool` spans are opened and closed from the same bus events the AG-UI stream is derived from, named after the AI Config key (not OpenCode's own tool id or the bridge's composed MCP name). diff --git a/packages/opencode-agents/agents.md b/packages/opencode-agents/agents.md index f6f22412..07ca3824 100644 --- a/packages/opencode-agents/agents.md +++ b/packages/opencode-agents/agents.md @@ -12,10 +12,12 @@ This package is Python-only. There is no JS counterpart in this slice. The spec | File | Role | | --- | --- | -| `handler.py` | Reducer, `create_opencode_handler`, `opencode_agents` | +| `handler.py` | Reducer, `create_opencode_handler`, `opencode_agents`, tool planning (`_plan_tools`, `_build_tool_ruleset`) | | `process.py` | Starts and stops `opencode serve` for one call | -| `server.py` | The HTTP session this process just started | -| `spans.py` | `invoke_agent` root, and the `chat` / `execute_tool` constructors the parity checks call | +| `server.py` | The HTTP session this process just started; `register_mcp_bridge`/`disconnect_mcp_bridge` | +| `spans.py` | `invoke_agent` root, the `chat` / `execute_tool` constructors, and `close_open_spans`/`abandon_open_spans` | +| `tool_bridge.py` | The per-call loopback MCP server a plain callable tool is served from | +| `builtins.py` | `NativeTool` sentinels for OpenCode's own built-ins (`OpenCodeBash`, etc.) | | `_version.py` | Package name and version sent with usage events | ## Dependencies @@ -38,9 +40,21 @@ A run opens one `invoke_agent` span and sets `launchdarkly.operation.type`. The Assistant text is emitted only after `message.updated` marks that message `role == "assistant"`. A part snapshot is the full text of that part so far; a later delta appends. A text delta without `partID` is dropped. +## Tool contract (Option C) + +`_plan_tools(config, tool_handlers)` is the single place that decides what one call offers: the intersection of `config["tools"]` and the keys present in `tool_handlers` (already wrapped by `tracking.py`, so a `NativeTool` entry carries `NATIVE_TOOL_KEY`). A `NativeTool` is allowed by its own `tool_name`. A plain callable is queued as a `_CallableEntry` and only gets a model-visible name once `_stream_impl` has started that call's `ToolBridge` and knows its generated MCP server name (`ld_`); `bridge_tool_name` composes the name OpenCode will actually report, matching `mcp/catalog.ts`'s own `sanitize` + `toolName`. + +`_build_tool_ruleset(offered_names)` is sent as `POST /session`'s `permission`: a leading `{"*": deny}` rule, then one `{name: allow}` rule per offered name. `Permission.merge` (OpenCode's own code) appends a session's `permission` after the agent's static `opencode.json` permission, and `Permission.disabled` resolves the *last* matching rule — so this ruleset always wins, and an empty `offered_names` still denies everything rather than falling back to the host's static catalog. + +The reducer (`_Reducer`) opens and closes `chat`/`execute_tool` spans from the same tool-part and message-info transitions the AG-UI events already come from, independent of what was actually offered — a span reports what OpenCode said happened, not what was requested. `tool_display_names` maps the name OpenCode reports (a native id or a bridge's composed name) back to the AI Config key for the span name; an unmapped name falls back to itself. A `NativeTool`'s own `$ld:ai:tool_call` tracking stub fires once, when its span opens (mirroring `claude-agents`'s `PreToolUse` hook); a bridged callable's tracked wrapper already fires it when the bridge invokes it, so the reducer does not fire it again. + +`open_tool_spans`/`open_model_spans` are owned by `_stream_impl`, not the reducer, specifically so `close_open_spans`/`abandon_open_spans` (registered in `test_cross_handler_parity.py`'s `HOOK_BASED_TOOL_HANDLERS`, this package's closest analogue to a vendor hook object) can find and end whatever is still open from the streaming `finally`, even if the reducer was never constructed (no raw bus event ever arrived). Keep both span-teardown signatures on one physical line — the parity suite's body extractor is indentation-based and mis-reads a wrapped signature whose closing `) -> None:` lands back at column 0, which is exactly what `ruff format` produces once a signature exceeds the line-length limit. + ## Common pitfalls -`provides_for` is `('*', 'agent')`, so a judge is routed here too. Each call starts its own process, so the judge prompt does not enter the agent's session and does not need `skip_judges=True`. The judge still runs as a full OpenCode agent, with the server's tools, in `cwd`. This package does not take a tool-less judge agent. +`provides_for` is `('*', 'agent')`, so a judge is routed here too. Each call starts its own process, so the judge prompt does not enter the agent's session and does not need `skip_judges=True`. Judges always call with `tool_handlers=None` (`judges.py`), so the offered set is empty and the judge gets a tool-less OpenCode agent — unlike `claude-agents`, which leaves Claude Code's built-ins on for a tool-less judge today; that inconsistency is open, not fixed here. + +The bridge (`tool_bridge.py`) answers in MCP streamable-HTTP "JSON response" mode only: no SSE, because this bridge never needs to push an unsolicited message to OpenCode. It runs on a background thread (`http.server`), not an asyncio server, to avoid taking on an ASGI dependency for one loopback endpoint; a `tools/call` crosses back onto the caller's event loop with `asyncio.run_coroutine_threadsafe`. None of this has been exercised against OpenCode's real MCP client (`@modelcontextprotocol/sdk`); `fake_opencode.py` only records the `POST /mcp` body, it never actually calls back into a registered bridge, except the one test in `test_tool_contract.py` that does so directly with a raw `httpx` request built from that recorded body. The reducer does not adopt the first session on the bus. Other sessions are dropped, including their `session.error`. `session.idle` for this session ends the run. diff --git a/packages/opencode-agents/api-report/launchdarkly_ai_opencode_agents.txt b/packages/opencode-agents/api-report/launchdarkly_ai_opencode_agents.txt index f5ad9a23..168a81df 100644 --- a/packages/opencode-agents/api-report/launchdarkly_ai_opencode_agents.txt +++ b/packages/opencode-agents/api-report/launchdarkly_ai_opencode_agents.txt @@ -1,4 +1,17 @@ # Public API of launchdarkly_ai_opencode_agents. Generated by `make api-report`; do not edit. create_opencode_handler(*, cwd: str | os.PathLike[str], config: str | os.PathLike[str], env: Mapping[str, str], capture_content: bool | None = None, model: ModelResolver | None = None, agent: str | None = None, command: Command = None) -> ProviderHandler opencode_agents(config_key: str, user_input: str, context: dict[str, Any], **kwargs: Any) -> Any +OpenCodeApplyPatch = NativeTool('apply_patch') +OpenCodeBash = NativeTool('bash') +OpenCodeEdit = NativeTool('edit') +OpenCodeGlob = NativeTool('glob') +OpenCodeGrep = NativeTool('grep') +OpenCodeList = NativeTool('list') +OpenCodeRead = NativeTool('read') OpenCodeSessionError(Exception) [class] +OpenCodeSkill = NativeTool('skill') +OpenCodeTask = NativeTool('task') +OpenCodeTodoWrite = NativeTool('todowrite') +OpenCodeWebFetch = NativeTool('webfetch') +OpenCodeWebSearch = NativeTool('websearch') +OpenCodeWrite = NativeTool('write') diff --git a/packages/opencode-agents/pyproject.toml b/packages/opencode-agents/pyproject.toml index ed0a1c7b..11995a80 100644 --- a/packages/opencode-agents/pyproject.toml +++ b/packages/opencode-agents/pyproject.toml @@ -3,7 +3,7 @@ name = "launchdarkly-ai-opencode-agents" version = "0.2.0" requires-python = ">=3.12" dependencies = [ - "launchdarkly-ai-server>=0.4.1", + "launchdarkly-ai-server>=0.4.0", "httpx>=0.27", "opentelemetry-api>=1.25", ] diff --git a/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/__init__.py b/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/__init__.py index 9b3ca199..ef35f836 100644 --- a/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/__init__.py +++ b/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/__init__.py @@ -6,10 +6,38 @@ from launchdarkly_ai_server import register_ai_sdk_package from ._version import PACKAGE_NAME, __version__ +from .builtins import ( + OpenCodeApplyPatch, + OpenCodeBash, + OpenCodeEdit, + OpenCodeGlob, + OpenCodeGrep, + OpenCodeList, + OpenCodeRead, + OpenCodeSkill, + OpenCodeTask, + OpenCodeTodoWrite, + OpenCodeWebFetch, + OpenCodeWebSearch, + OpenCodeWrite, +) from .handler import OpenCodeSessionError, create_opencode_handler, opencode_agents __all__ = [ + "OpenCodeApplyPatch", + "OpenCodeBash", + "OpenCodeEdit", + "OpenCodeGlob", + "OpenCodeGrep", + "OpenCodeList", + "OpenCodeRead", "OpenCodeSessionError", + "OpenCodeSkill", + "OpenCodeTask", + "OpenCodeTodoWrite", + "OpenCodeWebFetch", + "OpenCodeWebSearch", + "OpenCodeWrite", "create_opencode_handler", "opencode_agents", ] diff --git a/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/builtins.py b/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/builtins.py new file mode 100644 index 00000000..6c98998f --- /dev/null +++ b/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/builtins.py @@ -0,0 +1,43 @@ +"""OpenCode built-in tool sentinels. + +Place any of these as a value in your ``tool_handlers`` map to allow the +corresponding OpenCode built-in capability for one call. Every other +built-in, including one not named here, is denied for that call: an AI +Config variation with no tools gets a model with no tools at all. + +Ids verified against the OpenCode source at tag ``v1.18.4`` +(``packages/opencode/src/tool/*.ts``). ``OpenCodeList`` is the one +exception: it mirrors Vega's own ``_PINNABLE_BUILTIN_TOOLS`` pin +(``ai/vega-agentcore/vega_agentcore/opencode_config.py``) even though that +tag ships no ``list`` tool file. Passing a name OpenCode does not currently +serve is harmless — it simply never matches a tool the model is offered. + +Example:: + + from launchdarkly_ai_opencode_agents import OpenCodeBash, OpenCodeRead + from launchdarkly_ai_server import config + + response = await config( + key="my-agent", + handler=create_opencode_handler(cwd=..., config=..., env=...), + tool_handlers={"run-bash": OpenCodeBash, "read-file": OpenCodeRead}, + ).invoke(user_input, context) +""" + +from __future__ import annotations + +from launchdarkly_ai_server import NativeTool + +OpenCodeBash = NativeTool("bash") +OpenCodeRead = NativeTool("read") +OpenCodeEdit = NativeTool("edit") +OpenCodeWrite = NativeTool("write") +OpenCodeGlob = NativeTool("glob") +OpenCodeGrep = NativeTool("grep") +OpenCodeList = NativeTool("list") +OpenCodeTask = NativeTool("task") +OpenCodeSkill = NativeTool("skill") +OpenCodeTodoWrite = NativeTool("todowrite") +OpenCodeWebFetch = NativeTool("webfetch") +OpenCodeWebSearch = NativeTool("websearch") +OpenCodeApplyPatch = NativeTool("apply_patch") diff --git a/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/handler.py b/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/handler.py index 6a2e9886..02a9e027 100644 --- a/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/handler.py +++ b/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/handler.py @@ -11,11 +11,13 @@ import asyncio import os -from collections.abc import AsyncGenerator, AsyncIterable, Mapping +import secrets +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Mapping from contextlib import aclosing from typing import Any from launchdarkly_ai_server import ( + NATIVE_TOOL_KEY, AiConfigRep, ProviderHandler, SpanMessage, @@ -28,6 +30,7 @@ parse_template, set_input_content_attributes, set_output_content_attributes, + set_tool_call_content_attributes, text_message, ) from launchdarkly_ai_server.agui import TurnEmitter @@ -38,8 +41,28 @@ from ._version import PACKAGE_NAME, __version__ from .process import Command, OpenCodeServe, serve_from -from .server import ModelResolver, default_model, session_events -from .spans import fail_span, finish_root_span, mark_ok, start_root_span +from .server import ( + MESSAGE_TIMEOUT_SECONDS, + ModelResolver, + default_model, + disconnect_mcp_bridge, + register_mcp_bridge, + session_events, +) +from .spans import ( + abandon_open_spans, + close_open_spans, + fail_span, + finish_model_span, + finish_root_span, + mark_ok, + parent_context_of, + start_model_span, + start_root_span, + start_tool_span, + succeed_span, +) +from .tool_bridge import ToolBridge, bridge_tool_name #: The ``model.parameters`` keys forwarded into the ``POST /session/{id}/message`` body. None are. #: The message route takes no sampling, length, stop, seed, tool-choice or reasoning field: OpenCode @@ -78,11 +101,43 @@ class _Reducer: Events whose ``sessionID`` is not *session_id* are ignored. A missing ``sessionID`` is kept. The reducer does not adopt the first session it sees. A ``session.idle`` for this session ends the iterable. + + Also opens and closes this run's ``chat`` and ``execute_tool`` spans, at + the same two places text and tool events are already derived from the + bus: one ``chat`` span per assistant message id, one ``execute_tool`` + span per tool call id. *open_tool_spans* and *open_model_spans* are + owned by the caller, so a teardown outside this reducer (abandonment, + the run raising) can still find and close whatever is left open. + *tool_display_names* maps the name OpenCode reports for a tool part + (a native tool's own id, or a bridged tool's composed MCP name) back to + the AI Config key a span should show; an unmapped name is shown as-is. + *tool_native_stubs* holds the zero-arg ``$ld:ai:tool_call`` tracking stub + for a ``NativeTool`` entry, called once per call id when its span opens — + a bridged tool's own tracked wrapper already does this when the bridge + calls it, so it has no entry here. """ - def __init__(self, *, capture_content: bool, session_id: str) -> None: + def __init__( + self, + *, + capture_content: bool, + session_id: str, + config: AiConfigRep, + parent: Any, + open_tool_spans: dict[str, Any], + open_model_spans: dict[str, Any], + tool_display_names: Mapping[str, str], + tool_native_stubs: Mapping[str, Callable[[], None]], + ) -> None: self._emitter = TurnEmitter(capture_content=capture_content) + self._capture = capture_content self._session_id = session_id + self._config = config + self._parent = parent + self._tool_spans = open_tool_spans + self._model_spans = open_model_spans + self._tool_display_names = tool_display_names + self._tool_native_stubs = tool_native_stubs self._started: set[str] = set() self._text_message_id: str | None = None self._text_open = False @@ -356,6 +411,7 @@ def _emit_tool(self, part: Mapping[str, Any]) -> list[dict[str, Any]]: if status == "running": if "input" not in state or not name: return [] + self._open_tool_span(call_id, name, state.get("input")) events = self._emitter.tool_start(call_id, name) if any(event.get("type") == "TOOL_CALL_START" for event in events): self._started.add(call_id) @@ -363,6 +419,11 @@ def _emit_tool(self, part: Mapping[str, Any]) -> list[dict[str, Any]]: return events if status not in ("completed", "error"): return [] + if name: + # A call that skipped straight to completed/error without a + # recorded `running` state still needs its span opened once, + # with whatever input this status carries. + self._open_tool_span(call_id, name, state.get("input")) opened: list[dict[str, Any]] = [] if name: started = self._emitter.tool_start(call_id, name) @@ -371,11 +432,43 @@ def _emit_tool(self, part: Mapping[str, Any]) -> list[dict[str, Any]]: opened.extend(started) opened.extend(self._emitter.tool_end(call_id, arguments=state.get("input"))) if status == "error": + self._close_tool_span(call_id, failed=True, result=None) opened.extend(self._emitter.tool_result(call_id, failed=True)) else: - opened.extend(self._emitter.tool_result(call_id, state.get("output"))) + output = state.get("output") + self._close_tool_span(call_id, failed=False, result=output) + opened.extend(self._emitter.tool_result(call_id, output)) return opened + def _open_tool_span(self, call_id: str, provider_name: str, arguments: Any) -> None: + """Opens one ``execute_tool`` span for *call_id*, once. + + Also fires a ``NativeTool``'s tracking stub, the first and only time + its span opens — mirroring the Claude handler's ``PreToolUse`` hook, + which calls the same kind of stub when it opens its own tool span. + A bridged (callable) tool has no stub here: its tracked wrapper + already emits ``$ld:ai:tool_call`` when the bridge invokes it. + """ + if call_id in self._tool_spans: + return + stub = self._tool_native_stubs.get(provider_name) + if stub is not None: + stub() + display_name = self._tool_display_names.get(provider_name, provider_name) + span = start_tool_span(display_name, call_id, self._parent) + self._tool_spans[call_id] = span + set_tool_call_content_attributes(span, self._capture, arguments=arguments) + + def _close_tool_span(self, call_id: str, *, failed: bool, result: Any) -> None: + span = self._tool_spans.pop(call_id, None) + if span is None: + return + if failed: + fail_span(span, OpenCodeSessionError(f"tool call {call_id} failed")) + return + set_tool_call_content_attributes(span, self._capture, result=result) + succeed_span(span) + def _message(self, info: Mapping[str, Any]) -> list[dict[str, Any]]: message_id = info.get("id") role = info.get("role") @@ -385,6 +478,10 @@ def _message(self, info: Mapping[str, Any]) -> list[dict[str, Any]]: if isinstance(message_id, str) and isinstance(role, str): self._roles[message_id] = role if role == "assistant": + if message_id not in self._model_spans: + self._model_spans[message_id] = start_model_span( + self._config, self._parent + ) # A reserved id is not an open AG-UI message. Only text that # TurnEmitter has already yielded can take a parent or be split. late_completion = ( @@ -418,8 +515,17 @@ def _message(self, info: Mapping[str, Any]) -> list[dict[str, Any]]: if isinstance(tokens, Mapping): self._add_usage(tokens) self._usage_ids.add(message_id) + self._close_model_span(message_id, info.get("tokens")) return events + def _close_model_span( + self, message_id: str, tokens: Mapping[str, Any] | None + ) -> None: + span = self._model_spans.pop(message_id, None) + if span is None: + return + finish_model_span(span, self._config, _message_usage(tokens)) + def _add_usage(self, tokens: Mapping[str, Any]) -> None: self._saw_usage = True self._input += number_or_zero(tokens.get("input")) @@ -431,6 +537,25 @@ def _add_usage(self, tokens: Mapping[str, Any]) -> None: self._cache_write += number_or_zero(cache.get("write")) +def _message_usage(tokens: Any) -> SpanUsage: + """One assistant message's own token counts, as a ``chat`` span wants them. + + Separate from ``_Reducer._add_usage``, which accumulates the run's total + for the ``done`` event: a ``chat`` span reports the one message it + belongs to, not the run so far. + """ + if not isinstance(tokens, Mapping): + return SpanUsage() + cache = tokens.get("cache") + cache = cache if isinstance(cache, Mapping) else {} + return SpanUsage( + input=number_or_zero(tokens.get("input")), + output=number_or_zero(tokens.get("output")), + cache_read=number_or_zero(cache.get("read")), + cache_creation=number_or_zero(cache.get("write")), + ) + + def _string_id(value: object) -> str | None: if isinstance(value, str) and value: return value @@ -578,6 +703,88 @@ def opencode_agents( ).invoke(user_input, context, variables, history) +class _CallableEntry: + __slots__ = ("description", "fn", "parameters") + + def __init__( + self, description: str, parameters: Mapping[str, Any], fn: Callable[[Any], Any] + ) -> None: + self.description = description + self.parameters = parameters + self.fn = fn + + +class _ToolPlan: + """What one call offers OpenCode, derived from ``config["tools"]`` and ``tool_handlers``. + + *offered_names* and *tool_display_names* start populated with every + ``NativeTool`` entry. The handler appends a bridged (callable) tool's + composed MCP name to both once the bridge, if any, has started and that + name is known — a bridged tool's model-visible name depends on the + per-call bridge's generated server name, which does not exist yet when + this plan is built. + """ + + def __init__(self) -> None: + self.callables: dict[str, _CallableEntry] = {} + self.native_stubs: dict[str, Callable[[], None]] = {} + self.offered_names: list[str] = [] + self.tool_display_names: dict[str, str] = {} + + +def _plan_tools(config: AiConfigRep, tool_handlers: dict[str, Any] | None) -> _ToolPlan: + """Partitions the offered tools into native pass-throughs and bridged callables. + + "Offered" is the request-scoped authorization boundary: a config tool key + with no entry in *tool_handlers* is not served, and a *tool_handlers* + entry whose key the config does not declare is not served either. See + AGENTS.md's tool-handling contract and the matching rule already stated + for the JS SDK's handlers. + """ + plan = _ToolPlan() + config_tools = config.get("tools") if isinstance(config, dict) else None + config_tools = config_tools if isinstance(config_tools, dict) else {} + handlers = tool_handlers or {} + for key, tool_cfg in config_tools.items(): + stub = handlers.get(key) + if stub is None: + continue + native = getattr(stub, NATIVE_TOOL_KEY, None) + if native is not None: + provider_name = native.tool_name + plan.offered_names.append(provider_name) + plan.tool_display_names[provider_name] = key + plan.native_stubs[provider_name] = stub + continue + cfg = tool_cfg if isinstance(tool_cfg, dict) else {} + description = cfg.get("description") + parameters = cfg.get("parameters") + plan.callables[key] = _CallableEntry( + description=description if isinstance(description, str) else "", + parameters=parameters if isinstance(parameters, dict) else {}, + fn=stub, + ) + return plan + + +def _build_tool_ruleset(offered_names: list[str]) -> list[dict[str, str]]: + """OpenCode's session-level ``permission`` ruleset for one call's offered tools. + + The first rule denies every tool. ``Permission.disabled`` resolves the + *last* matching rule (``findLast`` over the agent's static permission + followed by this ruleset), so each later per-name allow rule takes + precedence for the tool it names and nothing else gains visibility by + accident — and this ruleset, being appended last, always wins over the + opencode.json agent config's own static permission. An empty + *offered_names* still sends the deny-all rule: a variation with no tools + gets a model with no tools, not the host's default catalog. + """ + ruleset = [{"permission": "*", "action": "deny", "pattern": "*"}] + for name in offered_names: + ruleset.append({"permission": name, "action": "allow", "pattern": "*"}) + return ruleset + + class OpenCodeHandler(ProviderHandler): """Handler whose async context owns one ``opencode serve`` process.""" @@ -646,13 +853,16 @@ async def _stream_impl( variables: dict[str, Any] | None, history: list[dict[str, Any]] | None = None, ) -> AsyncGenerator[dict[str, Any], None]: - del tool_handlers capture = effective_capture_content(self._capture_content_override) span = start_root_span(config, variables or {}) + parent = parent_context_of(span) ended: set[int] = set() + open_tool_spans: dict[str, Any] = {} + open_model_spans: dict[str, Any] = {} provider_id, model_id = self._model(config if isinstance(config, dict) else {}) system = _system_text(config, variables) prompt = _user_prompt(config, user_input, history, variables) + plan = _plan_tools(config, tool_handlers) # A context-owned process outlives this stream, so a caller timeout # that cancels the stream cannot cancel its shutdown. owned: OpenCodeServe | None = None @@ -667,6 +877,8 @@ async def _stream_impl( owned = self._serve.spawn() base_url = "" password = "" + bridge: ToolBridge | None = None + bridge_name: str | None = None try: set_input_content_attributes( span, @@ -677,6 +889,31 @@ async def _stream_impl( if owned is not None: base_url = await owned.start() password = owned.password + if plan.callables: + bridge = ToolBridge() + for key, entry in plan.callables.items(): + bridge.register( + key, + description=entry.description, + parameters=entry.parameters, + fn=entry.fn, + ) + await bridge.start() + run_id = secrets.token_hex(8) + bridge_name = f"ld_{run_id}" + await register_mcp_bridge( + base_url=base_url, + password=password, + name=bridge_name, + url=f"{bridge.base_url}/mcp/{run_id}", + token=bridge.token, + timeout_ms=int(MESSAGE_TIMEOUT_SECONDS * 1000), + ) + for key in plan.callables: + composed = bridge_tool_name(bridge_name, key) + plan.offered_names.append(composed) + plan.tool_display_names[composed] = key + permission = _build_tool_ruleset(plan.offered_names) created = _Created() async with aclosing( _reduce_session( @@ -692,6 +929,13 @@ async def _stream_impl( parameters=select_forwarded_parameters( model_parameters(config), _PROMPT_FORWARDED_KEYS ), + permission=permission, + config=config, + parent=parent, + open_tool_spans=open_tool_spans, + open_model_spans=open_model_spans, + tool_display_names=plan.tool_display_names, + tool_native_stubs=plan.native_stubs, ) ) as reduced: async for event in reduced: @@ -713,11 +957,19 @@ async def _stream_impl( done["usage"] = usage yield done except Exception as exc: + close_open_spans(open_tool_spans, open_model_spans, exc) fail_span(span, exc, ended) raise finally: + if bridge_name is not None: + await disconnect_mcp_bridge( + base_url=base_url, password=password, name=bridge_name + ) + if bridge is not None: + await bridge.stop() if owned is not None: await owned.stop() + abandon_open_spans(open_tool_spans, open_model_spans, ended) end_span_once(span, ended) async def _call_impl( @@ -764,6 +1016,13 @@ async def _reduce_session( agent: str | None, capture: bool, parameters: Mapping[str, Any], + permission: list[dict[str, str]] | None, + config: AiConfigRep, + parent: Any, + open_tool_spans: dict[str, Any], + open_model_spans: dict[str, Any], + tool_display_names: Mapping[str, str], + tool_native_stubs: Mapping[str, Callable[[], None]], ) -> AsyncGenerator[dict[str, Any], None]: """Reduce the session ``POST /session`` created. The reducer is unchanged.""" session_ids: list[str] = [] @@ -779,11 +1038,21 @@ async def _reduce_session( model_id=model_id, agent=agent, session_ids=session_ids, + permission=permission, ) ) as raw_events: async for raw in raw_events: if reducer is None: - reducer = _Reducer(capture_content=capture, session_id=session_ids[0]) + reducer = _Reducer( + capture_content=capture, + session_id=session_ids[0], + config=config, + parent=parent, + open_tool_spans=open_tool_spans, + open_model_spans=open_model_spans, + tool_display_names=tool_display_names, + tool_native_stubs=tool_native_stubs, + ) async with aclosing(reducer.reduce(_once(raw))) as reduced: async for event in reduced: yield event diff --git a/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/server.py b/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/server.py index b4a96e3b..71aae982 100644 --- a/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/server.py +++ b/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/server.py @@ -65,6 +65,7 @@ async def session_events( model_id: str, agent: str | None, session_ids: list[str], + permission: list[dict[str, str]] | None = None, ) -> AsyncGenerator[dict[str, Any], None]: """Open one session and yield its bus events. @@ -75,6 +76,13 @@ async def session_events( The post runs concurrently with the read. This session's ``session.idle`` ends the turn. A failed message POST raises without waiting for another bus frame. Leaving early cancels the post and aborts the session. + + *permission* is sent as the session's own ``permission`` ruleset + (``Session.CreateInput.permission``). ``Permission.merge`` appends it + after the agent's static config permission and ``Permission.disabled`` + resolves the *last* matching rule, so this always wins over whatever the + opencode.json agent config says: the catalog this call is offering is + authoritative, not the host's static default. """ http = httpx.AsyncClient( auth=(SERVER_USERNAME, password), @@ -91,7 +99,12 @@ async def session_events( timeout=httpx.Timeout(HTTP_TIMEOUT_SECONDS, read=None), ) as sse: sse.raise_for_status() - created = await http.post(f"{base}/session", timeout=HTTP_TIMEOUT_SECONDS) + session_body = {"permission": permission} if permission else None + created = await http.post( + f"{base}/session", + json=session_body, + timeout=HTTP_TIMEOUT_SECONDS, + ) if created.status_code >= 400: raise OpenCodeHTTPError(f"POST /session -> {created.status_code}") session_id = _session_id(created.json()) @@ -139,6 +152,53 @@ async def session_events( await http.aclose() +async def register_mcp_bridge( + *, base_url: str, password: str, name: str, url: str, token: str, timeout_ms: int +) -> None: + """``POST /mcp``: attaches a remote MCP server to the running instance. + + Per-call, not per-session: the route adds the server to the instance + (``MCP.add``), so this must run before the session's message is posted, + not before the session itself — there is no per-session MCP attachment. + """ + base = base_url.rstrip("/") + async with httpx.AsyncClient( + auth=(SERVER_USERNAME, password), timeout=HTTP_TIMEOUT_SECONDS + ) as http: + response = await http.post( + f"{base}/mcp", + json={ + "name": name, + "config": { + "type": "remote", + "url": url, + "headers": {"Authorization": f"Bearer {token}"}, + "oauth": False, + "timeout": timeout_ms, + }, + }, + ) + if response.status_code >= 400: + raise OpenCodeHTTPError(f"POST /mcp -> {response.status_code}") + + +async def disconnect_mcp_bridge(*, base_url: str, password: str, name: str) -> None: + """``POST /mcp/{name}/disconnect``, best-effort. + + Swallows every error: this runs from a ``finally``, after the turn is + already done, and the opencode process this call started (or shares) is + about to be told to stop regardless of whether this succeeds. + """ + base = base_url.rstrip("/") + try: + async with httpx.AsyncClient( + auth=(SERVER_USERNAME, password), timeout=HTTP_TIMEOUT_SECONDS + ) as http: + await http.post(f"{base}/mcp/{name}/disconnect") + except httpx.HTTPError: + pass + + async def _read_until_done( response: httpx.Response, post: asyncio.Task[httpx.Response], diff --git a/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/spans.py b/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/spans.py index 0b2aa2d2..a473c6a0 100644 --- a/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/spans.py +++ b/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/spans.py @@ -13,6 +13,7 @@ from launchdarkly_ai_server import ( AiConfigRep, SpanUsage, + end_span_once, set_ld_span_attributes, set_model_identity_attributes, set_usage_span_attributes, @@ -139,8 +140,42 @@ def fail_span(span: Any, error: BaseException, tracker: set[int] | None = None) span.record_exception(error) span.set_status(SpanStatusCode.ERROR, str(error)) if tracker is not None: - from launchdarkly_ai_server import end_span_once - end_span_once(span, tracker) else: span.end() + + +#: A call id (tool spans) or message id (model spans) to its still-open span. +Spans = dict[str, Any] + + +def close_open_spans(tools: Spans, models: Spans, error: BaseException) -> None: + """Fails every ``execute_tool`` and ``chat`` span still open. + + For the path where the run raises mid-stream: the reducer's own closing + logic never got to run for a call or a message that was still in flight, + so whatever is left in either dict would otherwise export as a span that + never ended. + """ + for span in list(tools.values()): + fail_span(span, error) + tools.clear() + for span in list(models.values()): + fail_span(span, error) + models.clear() + + +def abandon_open_spans(tools: Spans, models: Spans, ended: set[int]) -> None: + """Ends every ``execute_tool`` and ``chat`` span still open, for a + consumer that stopped reading. + + Nothing failed here: a consumer walking away is a normal thing to do, so + each span is left UNSET and marked ``launchdarkly.stream.abandoned``, + matching what the hook-based handlers do on the same path. + """ + for span in list(tools.values()): + end_span_once(span, ended, abandoned=True) + tools.clear() + for span in list(models.values()): + end_span_once(span, ended, abandoned=True) + models.clear() diff --git a/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/tool_bridge.py b/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/tool_bridge.py new file mode 100644 index 00000000..e0b27bed --- /dev/null +++ b/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/tool_bridge.py @@ -0,0 +1,278 @@ +"""A loopback MCP server for AI Config tools backed by plain callables. + +OpenCode has no route that accepts a tool definition alongside a message: the +message route's ``tools`` field is a deprecated boolean permission map, not a +place to describe a tool (see ``session/prompt.ts`` at the pinned OpenCode +tag). A plain callable in ``tool_handlers`` is therefore served from a +minimal MCP endpoint this module starts per call, which the handler attaches +to the running ``opencode serve`` instance with ``POST /mcp``. + +This endpoint answers in MCP streamable-HTTP "JSON response" mode: every +``tools/call`` gets exactly one JSON-RPC response body, with no +server-initiated SSE stream. That is a conforming subset of the transport +(a server with nothing to push to the client is not required to stream), and +it is all this bridge needs: OpenCode is the only client, and it never needs +an unsolicited message from this side. + +The stdlib has no asyncio HTTP server, and this package otherwise depends on +nothing beyond ``httpx`` and ``opentelemetry-api``. Bringing in an ASGI stack +for one loopback endpoint would be a much larger footprint than the endpoint +itself, so this runs on a background thread via ``http.server``, and crosses +back onto the event loop that is driving the call with +``asyncio.run_coroutine_threadsafe`` to invoke the tracked tool handler. + +Each call gets its own :class:`ToolBridge`, its own ephemeral port, and its +own bearer token, so two calls sharing one ``opencode serve`` process +(``async with handler``) can never reach each other's tools. +""" + +from __future__ import annotations + +import asyncio +import json +import re +import secrets +import threading +from collections.abc import Awaitable, Callable, Mapping +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any + +#: OpenCode's own MCP tool-name sanitizer (``mcp/catalog.ts``'s ``sanitize``). +#: Matched exactly so the composed name this module predicts is the name +#: OpenCode actually shows the model. +_SANITIZE_RE = re.compile(r"[^a-zA-Z0-9_-]") + +#: Generous on purpose. A served tool may call out (the broker, a database), +#: and the default MCP request timeout (5s, ``core/src/v1/config/mcp.ts``) is +#: far too short for that. ``register_mcp_bridge`` sends this as the +#: connection's own ``timeout``, in milliseconds. +CALL_TIMEOUT_SECONDS = 600.0 + +_PROTOCOL_VERSION = "2025-06-18" +_SERVER_NAME = "launchdarkly-ai-opencode-agents" + + +def sanitize(value: str) -> str: + """OpenCode's own MCP tool-name sanitizer, reproduced so names agree.""" + return _SANITIZE_RE.sub("_", value) + + +def bridge_tool_name(server_name: str, key: str) -> str: + """The name OpenCode shows the model for one bridged tool. + + Matches ``MCP.toolName`` (``sanitize(clientName) + "_" + sanitize(name)``) + so a permission rule built from this name actually matches what OpenCode + reports back on the tool part. + """ + return f"{sanitize(server_name)}_{sanitize(key)}" + + +ToolCallable = Callable[[Any], "Awaitable[Any] | Any"] + + +class _ToolEntry: + __slots__ = ("description", "fn", "parameters") + + def __init__( + self, description: str, parameters: Mapping[str, Any], fn: ToolCallable + ) -> None: + self.description = description + self.parameters = parameters + self.fn = fn + + +class ToolBridge: + """One loopback MCP endpoint, scoped to a single handler call. + + Register every tool before :meth:`start`, or afterward — ``tools/list`` + reads the live map either way. ``start`` binds an ephemeral port and + returns its base URL; ``stop`` closes the socket and joins the serving + thread. + """ + + def __init__(self) -> None: + self._tools: dict[str, _ToolEntry] = {} + self._token = secrets.token_urlsafe(32) + self._server: ThreadingHTTPServer | None = None + self._thread: threading.Thread | None = None + self._loop: asyncio.AbstractEventLoop | None = None + self._port = 0 + + @property + def token(self) -> str: + """The bearer token a caller must present. Generated once per bridge.""" + return self._token + + @property + def base_url(self) -> str: + if not self._port: + raise RuntimeError("ToolBridge has not started.") + return f"http://127.0.0.1:{self._port}" + + def register( + self, + key: str, + *, + description: str, + parameters: Mapping[str, Any], + fn: ToolCallable, + ) -> None: + """Serves *fn* under *key* as an MCP tool. *key* is the raw (unsanitized) name.""" + self._tools[key] = _ToolEntry(description, parameters, fn) + + async def start(self) -> str: + """Binds a loopback port and starts serving. Returns the base URL.""" + self._loop = asyncio.get_running_loop() + bridge = self + + class _Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, fmt: str, *args: Any) -> None: + return + + def do_GET(self) -> None: + # No server-initiated messages: this bridge never needs the + # SSE half of the streamable-HTTP transport. + self.send_error(405) + + def do_POST(self) -> None: + if not bridge._authorized(self.headers): + self.send_error(401) + return + length = int(self.headers.get("Content-Length", "0") or "0") + raw = self.rfile.read(length) if length else b"" + try: + message = json.loads(raw.decode("utf-8")) if raw else {} + except (ValueError, UnicodeDecodeError): + message = {} + response = bridge._dispatch( + message if isinstance(message, dict) else {} + ) + if response is None: + # A notification (no `id`): MCP streamable HTTP wants 202, no body. + self.send_response(202) + self.send_header("Content-Length", "0") + self.end_headers() + return + body = json.dumps(response).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + server = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) + server.daemon_threads = True + self._server = server + self._port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + self._thread = thread + thread.start() + return self.base_url + + async def stop(self) -> None: + server = self._server + self._server = None + if server is not None: + server.shutdown() + server.server_close() + thread = self._thread + self._thread = None + if thread is not None: + thread.join(timeout=5.0) + + def _authorized(self, headers: Any) -> bool: + header = headers.get("Authorization") + if not isinstance(header, str): + return False + return secrets.compare_digest(header, f"Bearer {self._token}") + + def _dispatch(self, message: Mapping[str, Any]) -> dict[str, Any] | None: + method = message.get("method") + msg_id = message.get("id") + if method == "initialize": + return { + "jsonrpc": "2.0", + "id": msg_id, + "result": { + "protocolVersion": _PROTOCOL_VERSION, + "capabilities": {"tools": {}}, + "serverInfo": {"name": _SERVER_NAME, "version": "1"}, + }, + } + if method is None or str(method).startswith("notifications/"): + # A notification carries no `id` and gets no JSON-RPC response. + return None + if msg_id is None: + return None + if method == "tools/list": + return { + "jsonrpc": "2.0", + "id": msg_id, + "result": { + "tools": [ + { + "name": name, + "description": entry.description, + "inputSchema": entry.parameters + or {"type": "object", "properties": {}}, + } + for name, entry in self._tools.items() + ] + }, + } + if method == "tools/call": + params = message.get("params") + params = params if isinstance(params, Mapping) else {} + name = params.get("name") + arguments = params.get("arguments") or {} + entry = self._tools.get(name) if isinstance(name, str) else None + if entry is None: + return { + "jsonrpc": "2.0", + "id": msg_id, + "error": {"code": -32602, "message": f"Unknown tool: {name!r}"}, + } + return self._call_tool(msg_id, entry, arguments) + return { + "jsonrpc": "2.0", + "id": msg_id, + "error": {"code": -32601, "message": f"Unknown method: {method!r}"}, + } + + def _call_tool( + self, msg_id: Any, entry: _ToolEntry, arguments: Any + ) -> dict[str, Any]: + loop = self._loop + if loop is None: + return { + "jsonrpc": "2.0", + "id": msg_id, + "result": { + "content": [{"type": "text", "text": "bridge is not running"}], + "isError": True, + }, + } + future = asyncio.run_coroutine_threadsafe(_invoke(entry.fn, arguments), loop) + try: + text, failed = future.result(timeout=CALL_TIMEOUT_SECONDS) + except Exception as exc: # the handler raised, or the call timed out + text, failed = str(exc), True + return { + "jsonrpc": "2.0", + "id": msg_id, + "result": {"content": [{"type": "text", "text": text}], "isError": failed}, + } + + +async def _invoke(fn: ToolCallable, arguments: Any) -> tuple[str, bool]: + from launchdarkly_ai_server import tool_result_string + + try: + result = fn(arguments) + if asyncio.iscoroutine(result): + result = await result + return tool_result_string(result), False + except Exception as exc: + return str(exc), True diff --git a/packages/opencode-agents/tests/fake_opencode.py b/packages/opencode-agents/tests/fake_opencode.py index 202ecaf3..b4d0a095 100644 --- a/packages/opencode-agents/tests/fake_opencode.py +++ b/packages/opencode-agents/tests/fake_opencode.py @@ -84,9 +84,12 @@ def __init__(self, spec: dict[str, Any], state_dir: str) -> None: self.spec = spec self.state_dir = state_dir self.sessions: list[str] = [] + self.session_bodies: list[dict[str, Any]] = [] self.messages: list[dict[str, Any]] = [] self.aborts: list[str] = [] self.requests: list[dict[str, str]] = [] + self.mcp_calls: list[dict[str, Any]] = [] + self.mcp_disconnects: list[str] = [] self.second_read = threading.Event() self.release = threading.Event() self._lock = threading.Lock() @@ -120,9 +123,12 @@ def _write(self) -> None: "username": os.environ.get("OPENCODE_SERVER_USERNAME", ""), "password": os.environ.get("OPENCODE_SERVER_PASSWORD", ""), "sessions": self.sessions, + "session_bodies": self.session_bodies, "messages": self.messages, "aborts": self.aborts, "requests": self.requests, + "mcp_calls": self.mcp_calls, + "mcp_disconnects": self.mcp_disconnects, } path = os.path.join(self.state_dir, f"{os.getpid()}.json") temporary = path + ".tmp" @@ -228,8 +234,20 @@ def serve_post(self, path: str, body: bytes) -> tuple[int, dict[str, Any]]: "ses" if not self.sessions else f"ses-{len(self.sessions) + 1}" ) self.sessions.append(session_id) + self.session_bodies.append(json.loads(body.decode()) if body else {}) self._write() return 200, {"id": session_id} + if path == "/mcp": + parsed = json.loads(body.decode()) if body else {} + self.mcp_calls.append(parsed) + self._write() + name = parsed.get("name", "") + return 200, {"status": {name: {"status": "connected"}}} + if path.endswith("/disconnect") and path.startswith("/mcp/"): + name = path[len("/mcp/") : -len("/disconnect")] + self.mcp_disconnects.append(name) + self._write() + return 200, True if path.endswith("/message"): self.messages.append(json.loads(body.decode()) if body else {}) self._write() diff --git a/packages/opencode-agents/tests/test_handler.py b/packages/opencode-agents/tests/test_handler.py index cee6de67..d8b46003 100644 --- a/packages/opencode-agents/tests/test_handler.py +++ b/packages/opencode-agents/tests/test_handler.py @@ -209,6 +209,27 @@ def requests(self) -> list[dict[str, str]]: found.extend(record.get("requests") or []) return found + @property + def session_bodies(self) -> list[dict[str, Any]]: + found: list[dict[str, Any]] = [] + for record in self._records(): + found.extend(record.get("session_bodies") or []) + return found + + @property + def mcp_calls(self) -> list[dict[str, Any]]: + found: list[dict[str, Any]] = [] + for record in self._records(): + found.extend(record.get("mcp_calls") or []) + return found + + @property + def mcp_disconnects(self) -> list[str]: + found: list[str] = [] + for record in self._records(): + found.extend(record.get("mcp_disconnects") or []) + return found + _scripts: list[Script] = [] @@ -288,7 +309,20 @@ def test_package_exports_the_handler(self) -> None: import launchdarkly_ai_opencode_agents as package assert set(package.__all__) == { + "OpenCodeApplyPatch", + "OpenCodeBash", + "OpenCodeEdit", + "OpenCodeGlob", + "OpenCodeGrep", + "OpenCodeList", + "OpenCodeRead", "OpenCodeSessionError", + "OpenCodeSkill", + "OpenCodeTask", + "OpenCodeTodoWrite", + "OpenCodeWebFetch", + "OpenCodeWebSearch", + "OpenCodeWrite", "create_opencode_handler", "opencode_agents", } diff --git a/packages/opencode-agents/tests/test_tool_bridge.py b/packages/opencode-agents/tests/test_tool_bridge.py new file mode 100644 index 00000000..0c7bfb65 --- /dev/null +++ b/packages/opencode-agents/tests/test_tool_bridge.py @@ -0,0 +1,253 @@ +"""The loopback MCP bridge: JSON-RPC behavior, independent of ``opencode serve``. + +OpenCode's own MCP client is not available in this test environment, so these +tests exercise the bridge's HTTP surface directly with ``httpx``, the same way +a conforming MCP streamable-HTTP client would: ``initialize``, then +``notifications/initialized``, then ``tools/list`` and ``tools/call``. +""" + +from __future__ import annotations + +import json +from typing import Any + +import httpx + +from launchdarkly_ai_opencode_agents.tool_bridge import ( + ToolBridge, + bridge_tool_name, + sanitize, +) + + +def test_sanitize_matches_opencodes_own_regex() -> None: + # mcp/catalog.ts: `value.replace(/[^a-zA-Z0-9_-]/g, "_")`. + assert sanitize("get weather!") == "get_weather_" + assert sanitize("already-ok_Name9") == "already-ok_Name9" + + +def test_bridge_tool_name_composes_server_and_key() -> None: + assert bridge_tool_name("ld_abc123", "get weather") == "ld_abc123_get_weather" + + +class TestToolBridgeProtocol: + async def _post( + self, base_url: str, token: str, message: dict[str, Any] | None + ) -> httpx.Response: + async with httpx.AsyncClient() as http: + return await http.post( + f"{base_url}/mcp/run-1", + json=message, + headers={"Authorization": f"Bearer {token}"}, + ) + + async def test_rejects_a_missing_or_wrong_token(self) -> None: + bridge = ToolBridge() + await bridge.start() + try: + async with httpx.AsyncClient() as http: + response = await http.post(f"{bridge.base_url}/mcp/run-1", json={}) + assert response.status_code == 401 + response = await self._post(bridge.base_url, "wrong-token", {}) + assert response.status_code == 401 + finally: + await bridge.stop() + + async def test_initialize_reports_tool_capability(self) -> None: + bridge = ToolBridge() + await bridge.start() + try: + response = await self._post( + bridge.base_url, + bridge.token, + {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}, + ) + body = response.json() + assert body["id"] == 1 + assert body["result"]["capabilities"] == {"tools": {}} + finally: + await bridge.stop() + + async def test_a_notification_gets_202_and_no_body(self) -> None: + bridge = ToolBridge() + await bridge.start() + try: + response = await self._post( + bridge.base_url, + bridge.token, + {"jsonrpc": "2.0", "method": "notifications/initialized"}, + ) + assert response.status_code == 202 + assert response.content == b"" + finally: + await bridge.stop() + + async def test_tools_list_reports_the_registered_schema(self) -> None: + bridge = ToolBridge() + bridge.register( + "get_weather", + description="Look up the weather", + parameters={"type": "object", "properties": {"city": {"type": "string"}}}, + fn=lambda args: f"sunny in {args['city']}", + ) + await bridge.start() + try: + response = await self._post( + bridge.base_url, + bridge.token, + {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}}, + ) + tools = response.json()["result"]["tools"] + assert len(tools) == 1 + assert tools[0]["name"] == "get_weather" + assert tools[0]["description"] == "Look up the weather" + assert tools[0]["inputSchema"]["properties"]["city"]["type"] == "string" + finally: + await bridge.stop() + + async def test_tools_call_runs_a_sync_handler_and_wraps_the_result(self) -> None: + bridge = ToolBridge() + bridge.register( + "get_weather", + description="", + parameters={}, + fn=lambda args: f"sunny in {args['city']}", + ) + await bridge.start() + try: + response = await self._post( + bridge.base_url, + bridge.token, + { + "jsonrpc": "2.0", + "id": 3, + "method": "tools/call", + "params": {"name": "get_weather", "arguments": {"city": "Oslo"}}, + }, + ) + result = response.json()["result"] + assert result["isError"] is False + assert result["content"] == [{"type": "text", "text": "sunny in Oslo"}] + finally: + await bridge.stop() + + async def test_tools_call_awaits_an_async_handler(self) -> None: + async def handler(args: dict[str, Any]) -> dict[str, Any]: + return {"ok": True, "echo": args} + + bridge = ToolBridge() + bridge.register("echo", description="", parameters={}, fn=handler) + await bridge.start() + try: + response = await self._post( + bridge.base_url, + bridge.token, + { + "jsonrpc": "2.0", + "id": 4, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"x": 1}}, + }, + ) + result = response.json()["result"] + assert result["isError"] is False + # Non-string results are JSON-encoded (tool_result_string). + assert json.loads(result["content"][0]["text"]) == { + "ok": True, + "echo": {"x": 1}, + } + finally: + await bridge.stop() + + async def test_a_raising_handler_comes_back_as_an_mcp_tool_error(self) -> None: + def boom(args: dict[str, Any]) -> str: + raise ValueError("no such city") + + bridge = ToolBridge() + bridge.register("get_weather", description="", parameters={}, fn=boom) + await bridge.start() + try: + response = await self._post( + bridge.base_url, + bridge.token, + { + "jsonrpc": "2.0", + "id": 5, + "method": "tools/call", + "params": {"name": "get_weather", "arguments": {}}, + }, + ) + result = response.json()["result"] + # A tool failure is an MCP-level `isError`, not a JSON-RPC protocol + # error: OpenCode is expected to surface it as a normal completed + # call whose part status is `error`, not reject the call outright. + assert result["isError"] is True + assert "no such city" in result["content"][0]["text"] + finally: + await bridge.stop() + + async def test_calling_an_unregistered_tool_is_a_protocol_error(self) -> None: + bridge = ToolBridge() + await bridge.start() + try: + response = await self._post( + bridge.base_url, + bridge.token, + { + "jsonrpc": "2.0", + "id": 6, + "method": "tools/call", + "params": {"name": "nope", "arguments": {}}, + }, + ) + body = response.json() + assert "error" in body + assert body["error"]["code"] == -32602 + finally: + await bridge.stop() + + async def test_an_unknown_method_is_a_protocol_error(self) -> None: + bridge = ToolBridge() + await bridge.start() + try: + response = await self._post( + bridge.base_url, + bridge.token, + {"jsonrpc": "2.0", "id": 7, "method": "resources/list"}, + ) + body = response.json() + assert body["error"]["code"] == -32601 + finally: + await bridge.stop() + + async def test_get_is_not_offered(self) -> None: + # This bridge never pushes an unsolicited message, so it does not + # implement the SSE half of the streamable-HTTP transport. + bridge = ToolBridge() + await bridge.start() + try: + async with httpx.AsyncClient() as http: + response = await http.get( + f"{bridge.base_url}/mcp/run-1", + headers={"Authorization": f"Bearer {bridge.token}"}, + ) + assert response.status_code == 405 + finally: + await bridge.stop() + + async def test_two_bridges_get_different_tokens_and_ports(self) -> None: + a = ToolBridge() + b = ToolBridge() + await a.start() + await b.start() + try: + assert a.token != b.token + assert a.base_url != b.base_url + finally: + await a.stop() + await b.stop() + + +async def test_stop_before_start_does_not_raise() -> None: + bridge = ToolBridge() + await bridge.stop() diff --git a/packages/opencode-agents/tests/test_tool_contract.py b/packages/opencode-agents/tests/test_tool_contract.py new file mode 100644 index 00000000..d51ae2b5 --- /dev/null +++ b/packages/opencode-agents/tests/test_tool_contract.py @@ -0,0 +1,670 @@ +"""The OpenCode handler honors AI Config tools (Option C of the tool-contract design). + +A variation's ``tools``, filtered to the keys the caller also passed in +``tool_handlers``, is the catalog OpenCode is allowed to show the model for +one call. A ``NativeTool`` is allowed by name. A plain callable is served +from this call's own loopback MCP bridge, registered on the running +``opencode serve`` with ``POST /mcp`` and torn down with +``POST /mcp/{name}/disconnect``. + +``fake_opencode.py`` never actually calls back into a registered MCP server: +it only records the ``POST /mcp`` body and plays back a scripted event list. +These tests verify the handler's side of that contract (the ruleset sent, +the registration body, the teardown call, and the spans and tracking a +scripted tool part produces) against the fake. The bridge's own MCP +protocol behavior is tested directly, with no fake opencode involved, in +``test_tool_bridge.py``; nothing here has been run against a real OpenCode +MCP client. +""" + +from __future__ import annotations + +import json +import os +import shutil +import sys +import tempfile +from collections.abc import Iterator +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +import launchdarkly_ai_server.lifecycle as lifecycle_module +from launchdarkly_ai_opencode_agents import ( + OpenCodeBash, + OpenCodeWebFetch, + create_opencode_handler, +) +from launchdarkly_ai_opencode_agents.handler import _build_tool_ruleset, _plan_tools +from launchdarkly_ai_server import NativeTool, wrap_tool_handlers + +CONTEXT = {"kind": "user", "key": "user-1"} +_TRACK_DATA: dict[str, Any] = { + "runId": "run-1", + "configKey": "cfg", + "variationKey": "v1", + "version": 1, + "modelName": "m", + "providerName": "p", +} + + +def _wrapped(handlers: dict[str, Any]) -> dict[str, Any]: + """``tool_handlers`` reach a handler already wrapped by ``tracking.py`` -- + a ``NativeTool`` becomes a tracking stub carrying ``NATIVE_TOOL_KEY``, and + a plain callable becomes an async wrapper. Calling the handler's + ``.stream()`` directly, as these tests do, bypasses that layer, so each + test wraps its own map the same way ``execute_and_stream`` would. + """ + return wrap_tool_handlers(handlers, CONTEXT, _TRACK_DATA) + + +# Duplicated from test_handler.py rather than imported: this directory has no +# `__init__.py`, pytest collects it with `--import-mode=importlib`, and a +# root-level `tests/` package already exists for the repo's own shared test +# utilities, so `tests.test_handler` does not resolve to this package's +# `test_handler.py`. `fake_opencode.py` is only ever invoked as a subprocess +# (never imported), so it is reused as-is, by path. + +_FAKE = Path(__file__).with_name("fake_opencode.py") + +VARIATION: dict[str, Any] = { + "model": {"name": "test-model"}, + "provider": {"name": "Acme"}, + "instructions": "Be brief.", + "_ldMeta": {"enabled": True, "variationKey": "v1", "version": 1, "mode": "agent"}, +} + + +def bus(kind: str, **properties: Any) -> dict[str, Any]: + return {"type": kind, "properties": properties} + + +def tool_part( + call_id: str, + status: str, + *, + message_id: str = "m1", + name: str = "bash", + raw_input: Any = None, + output: Any = None, + part_id: str | None = None, +) -> dict[str, Any]: + state: dict[str, Any] = {"status": status} + if raw_input is not None: + state["input"] = raw_input + if output is not None: + state["output"] = output + return bus( + "message.part.updated", + part={ + "id": part_id or f"part-{call_id}", + "messageID": message_id, + "type": "tool", + "callID": call_id, + "tool": name, + "state": state, + }, + ) + + +def assistant_open(message_id: str) -> dict[str, Any]: + return bus( + "message.updated", + info={"id": message_id, "role": "assistant", "time": {"created": 1}}, + ) + + +def assistant_completed( + message_id: str, *, tokens: dict[str, Any] | None = None +) -> dict[str, Any]: + info: dict[str, Any] = { + "id": message_id, + "role": "assistant", + "time": {"created": 1, "completed": 2}, + } + if tokens is not None: + info["tokens"] = tokens + return bus("message.updated", info=info) + + +class Script: + """A fake ``opencode`` executable. The handler starts and stops it.""" + + def __init__(self, events: list[dict[str, Any]]) -> None: + self.root = Path(tempfile.mkdtemp(prefix="oc-fake-")) + self.cwd = self.root / "work" + self.cwd.mkdir() + self.config = self.root / "opencode.json" + self.config.write_text("{}", encoding="utf-8") + self.state = self.root / "state" + self.state.mkdir() + self.spec_path = self.root / "spec.json" + self.spec_path.write_text( + json.dumps({"mode": "script", "events": events}), encoding="utf-8" + ) + _scripts.append(self) + + def handler_args(self) -> dict[str, Any]: + return { + "cwd": self.cwd, + "config": self.config, + "env": { + "LD_OPENCODE_FAKE_DIR": str(self.state), + "LD_OPENCODE_FAKE_SPEC": str(self.spec_path), + }, + "command": [sys.executable, str(_FAKE)], + } + + def _records(self) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + for path in sorted( + self.state.glob("*.json"), key=lambda item: item.stat().st_mtime + ): + records.append(json.loads(path.read_text(encoding="utf-8"))) + return records + + @property + def pids(self) -> list[int]: + return [int(record["pid"]) for record in self._records() if "pid" in record] + + @property + def session_bodies(self) -> list[dict[str, Any]]: + found: list[dict[str, Any]] = [] + for record in self._records(): + found.extend(record.get("session_bodies") or []) + return found + + @property + def mcp_calls(self) -> list[dict[str, Any]]: + found: list[dict[str, Any]] = [] + for record in self._records(): + found.extend(record.get("mcp_calls") or []) + return found + + @property + def mcp_disconnects(self) -> list[str]: + found: list[str] = [] + for record in self._records(): + found.extend(record.get("mcp_disconnects") or []) + return found + + +_scripts: list[Script] = [] + + +def playing(events: list[dict[str, Any]], **kwargs: Any) -> tuple[Any, Script]: + script = Script(events) + handler = create_opencode_handler(**script.handler_args(), **kwargs) + return handler, script + + +async def _drain(stream: Any) -> None: + async for _event in stream: + pass + + +@pytest.fixture(autouse=True) +def _reset_sdk() -> Iterator[None]: + lifecycle_module._reset_for_testing() + yield + lifecycle_module._reset_for_testing() + for script in _scripts: + for pid in script.pids: + try: + os.kill(pid, 15) + except OSError: + pass + shutil.rmtree(script.root, ignore_errors=True) + _scripts.clear() + + +@pytest.fixture +def mock_ld_client() -> Iterator[MagicMock]: + client = MagicMock() + client.track = MagicMock() + client.flush = AsyncMock() + client.close = AsyncMock() + client.variation = AsyncMock(return_value=VARIATION) + lifecycle_module._set_client_for_testing(client) + yield client + lifecycle_module._reset_for_testing() + + +_CONFIG_WITH_TOOLS: dict[str, Any] = { + "instructions": "Be brief.", + "model": {"name": "m"}, + "provider": {"name": "Acme"}, + "tools": { + "run_shell": { + "name": "run_shell", + "type": "function", + "description": "Run a shell command", + }, + "fetch_url": { + "name": "fetch_url", + "type": "function", + "description": "Fetch a URL", + }, + "get_weather": { + "name": "get_weather", + "type": "function", + "description": "Look up the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + "unused": { + "name": "unused", + "type": "function", + "description": "No handler is registered for this key", + }, + }, +} + + +class TestPlanTools: + """``_plan_tools`` and ``_build_tool_ruleset`` as pure functions.""" + + def test_offered_is_the_intersection_of_config_tools_and_tool_handlers( + self, + ) -> None: + plan = _plan_tools( + _CONFIG_WITH_TOOLS, + _wrapped( + { + "run_shell": OpenCodeBash, + # Not a declared config tool: must not be offered. + "extra": OpenCodeWebFetch, + } + ), + ) + assert plan.offered_names == ["bash"] + assert plan.callables == {} + + def test_a_config_tool_with_no_handler_is_not_offered(self) -> None: + plan = _plan_tools(_CONFIG_WITH_TOOLS, _wrapped({"run_shell": OpenCodeBash})) + assert "unused" not in plan.tool_display_names.values() + assert "fetch_url" not in plan.callables + assert "get_weather" not in plan.callables + + def test_native_tool_is_offered_by_its_provider_name(self) -> None: + plan = _plan_tools(_CONFIG_WITH_TOOLS, _wrapped({"run_shell": OpenCodeBash})) + assert plan.offered_names == ["bash"] + assert plan.tool_display_names["bash"] == "run_shell" + assert "bash" in plan.native_stubs + + def test_a_custom_named_native_tool_keeps_its_own_provider_name(self) -> None: + custom = NativeTool("observability_query_logs") + plan = _plan_tools( + {**_CONFIG_WITH_TOOLS, "tools": {"query_logs": {"name": "query_logs"}}}, + _wrapped({"query_logs": custom}), + ) + assert plan.offered_names == ["observability_query_logs"] + assert plan.tool_display_names["observability_query_logs"] == "query_logs" + + def test_a_plain_callable_becomes_a_pending_bridge_entry(self) -> None: + async def get_weather(args: dict[str, Any]) -> str: + return f"sunny in {args['city']}" + + plan = _plan_tools(_CONFIG_WITH_TOOLS, _wrapped({"get_weather": get_weather})) + assert ( + plan.offered_names == [] + ) # not named yet: the bridge has no name until it starts + assert set(plan.callables) == {"get_weather"} + entry = plan.callables["get_weather"] + assert entry.description == "Look up the weather" + assert entry.parameters["properties"]["city"]["type"] == "string" + + def test_no_tools_at_all_plans_nothing(self) -> None: + plan = _plan_tools(_CONFIG_WITH_TOOLS, None) + assert plan.offered_names == [] + assert plan.callables == {} + plan2 = _plan_tools( + {"instructions": "hi"}, _wrapped({"run_shell": OpenCodeBash}) + ) + assert plan2.offered_names == [] + + def test_ruleset_always_denies_before_allowing(self) -> None: + ruleset = _build_tool_ruleset(["bash", "ld_abc_get_weather"]) + assert ruleset[0] == {"permission": "*", "action": "deny", "pattern": "*"} + assert ruleset[1:] == [ + {"permission": "bash", "action": "allow", "pattern": "*"}, + {"permission": "ld_abc_get_weather", "action": "allow", "pattern": "*"}, + ] + + def test_an_empty_offered_set_still_denies_everything(self) -> None: + assert _build_tool_ruleset([]) == [ + {"permission": "*", "action": "deny", "pattern": "*"} + ] + + +class TestSessionPermission: + """What the handler actually sends OpenCode, through the fake server.""" + + async def test_no_tools_sends_the_deny_all_ruleset(self) -> None: + handler, script = playing([assistant_open("m1"), assistant_completed("m1")]) + await _drain(await handler.stream({}, "hi", None, None)) + assert script.session_bodies[-1]["permission"] == [ + {"permission": "*", "action": "deny", "pattern": "*"} + ] + + async def test_a_native_tool_is_allowed_by_its_provider_name(self) -> None: + handler, script = playing([assistant_open("m1"), assistant_completed("m1")]) + await _drain( + await handler.stream( + _CONFIG_WITH_TOOLS, "hi", _wrapped({"run_shell": OpenCodeBash}), None + ) + ) + permission = script.session_bodies[-1]["permission"] + assert permission[0] == {"permission": "*", "action": "deny", "pattern": "*"} + assert {"permission": "bash", "action": "allow", "pattern": "*"} in permission + # Only the offered tool is allowed -- not every config tool. + assert len(permission) == 2 + + async def test_a_callable_tool_registers_an_mcp_bridge_then_disconnects( + self, + ) -> None: + async def get_weather(args: dict[str, Any]) -> str: + return f"sunny in {args['city']}" + + handler, script = playing([assistant_open("m1"), assistant_completed("m1")]) + await _drain( + await handler.stream( + _CONFIG_WITH_TOOLS, "hi", _wrapped({"get_weather": get_weather}), None + ) + ) + assert len(script.mcp_calls) == 1 + call = script.mcp_calls[0] + name = call["name"] + assert name.startswith("ld_") + assert call["config"]["type"] == "remote" + assert call["config"]["url"].startswith("http://127.0.0.1:") + assert call["config"]["url"].endswith(f"/mcp/{name.removeprefix('ld_')}") + assert call["config"]["headers"]["Authorization"].startswith("Bearer ") + assert call["config"]["oauth"] is False + assert call["config"]["timeout"] >= 600_000 + + permission = script.session_bodies[-1]["permission"] + allowed = { + rule["permission"] for rule in permission if rule["action"] == "allow" + } + assert allowed == {f"{name}_get_weather"} + + assert script.mcp_disconnects == [name] + + async def test_native_and_callable_tools_can_be_offered_together(self) -> None: + async def get_weather(args: dict[str, Any]) -> str: + return "sunny" + + handler, script = playing([assistant_open("m1"), assistant_completed("m1")]) + await _drain( + await handler.stream( + _CONFIG_WITH_TOOLS, + "hi", + _wrapped({"run_shell": OpenCodeBash, "get_weather": get_weather}), + None, + ) + ) + permission = script.session_bodies[-1]["permission"] + allow_names = { + rule["permission"] for rule in permission if rule["action"] == "allow" + } + assert "bash" in allow_names + assert any(name.endswith("_get_weather") for name in allow_names) + assert len(script.mcp_calls) == 1 + + async def test_two_calls_on_one_shared_process_get_different_bridge_names( + self, + ) -> None: + async def get_weather(args: dict[str, Any]) -> str: + return "sunny" + + script = Script([assistant_open("m1"), assistant_completed("m1")]) + from launchdarkly_ai_opencode_agents import create_opencode_handler + + handler = create_opencode_handler(**script.handler_args()) + async with handler: + await _drain( + await handler.stream( + _CONFIG_WITH_TOOLS, + "hi", + _wrapped({"get_weather": get_weather}), + None, + ) + ) + await _drain( + await handler.stream( + _CONFIG_WITH_TOOLS, + "hi", + _wrapped({"get_weather": get_weather}), + None, + ) + ) + names = [call["name"] for call in script.mcp_calls] + assert len(names) == 2 + assert names[0] != names[1] + assert script.mcp_disconnects == names + + +class TestToolBridgeEndToEnd: + """Drives a registered bridge exactly as the fake session recorded it. + + The bridge is registered with ``POST /mcp`` before the session is even + created, so pulling just the first yielded AG-UI event (the tool call + the scripted session reports) is enough to know registration already + happened and the bridge is still alive: the handler's ``finally`` that + stops it has not run yet, because the generator is suspended mid-stream. + """ + + async def test_the_registered_bridge_answers_a_real_tools_call( + self, mock_ld_client: MagicMock + ) -> None: + async def get_weather(args: dict[str, Any]) -> str: + return f"sunny in {args['city']}" + + handler, script = playing( + [ + assistant_open("m1"), + tool_part("c1", "running", name="bash", raw_input={"command": "ls"}), + ] + ) + agen = ( + await handler.stream( + _CONFIG_WITH_TOOLS, "hi", _wrapped({"get_weather": get_weather}), None + ) + ).__aiter__() + first = await agen.__anext__() + assert first["type"] == "TOOL_CALL_START" + call = script.mcp_calls[0] + async with httpx.AsyncClient() as http: + response = await http.post( + call["config"]["url"], + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + # The bridge's own tool name is the bare AI Config key, not + # the `_` name OpenCode composes for the + # model: that composition is applied only to what the + # model sees, never sent back in `tools/call`. + "params": {"name": "get_weather", "arguments": {"city": "Oslo"}}, + }, + headers=call["config"]["headers"], + ) + result = response.json()["result"] + assert result["isError"] is False, result + assert result["content"] == [{"type": "text", "text": "sunny in Oslo"}] + await agen.aclose() + + +class TestToolSpans: + """``chat`` and ``execute_tool`` spans, opened from the same bus events as the AG-UI stream.""" + + @staticmethod + def _tracer(monkeypatch: pytest.MonkeyPatch): + from launchdarkly_ai_opencode_agents import spans as span_module + + class _Span: + def __init__(self, name: str, context: Any = None) -> None: + self.name = name + self.context = context + self.attributes: dict[str, Any] = {} + self.status: Any = None + self.status_description: str | None = None + self.exceptions: list[BaseException] = [] + self.ended = 0 + + def set_attribute(self, key: str, value: Any) -> None: + self.attributes[key] = value + + def add_event( + self, name: str, attributes: dict[str, Any] | None = None + ) -> None: + return None + + def set_status(self, code: Any, description: str | None = None) -> None: + self.status = code + self.status_description = description + + def record_exception(self, exc: BaseException) -> None: + self.exceptions.append(exc) + + def end(self, end_time: int | None = None) -> None: + self.ended += 1 + + class _Tracer: + def __init__(self) -> None: + self.spans: list[_Span] = [] + + def get_tracer(self, name: str) -> _Tracer: + return self + + def start_span( + self, name: str, context: Any = None, start_time: int | None = None + ) -> _Span: + span = _Span(name, context) + self.spans.append(span) + return span + + def set_span_in_context(self, span: _Span) -> Any: + return ("context-of", span) + + tracer = _Tracer() + monkeypatch.setattr(span_module, "trace", tracer) + return tracer + + async def test_a_native_tool_call_opens_execute_tool_under_the_config_key( + self, monkeypatch: pytest.MonkeyPatch, mock_ld_client: MagicMock + ) -> None: + tracer = self._tracer(monkeypatch) + handler, _script = playing( + [ + assistant_open("m1"), + tool_part("c1", "running", name="bash", raw_input={"command": "ls"}), + tool_part( + "c1", + "completed", + name="bash", + raw_input={"command": "ls"}, + output="a.txt", + ), + assistant_completed("m1"), + ] + ) + await _drain( + await handler.stream( + _CONFIG_WITH_TOOLS, "hi", _wrapped({"run_shell": OpenCodeBash}), None + ) + ) + tool_spans = [s for s in tracer.spans if s.name.startswith("execute_tool")] + assert len(tool_spans) == 1 + span = tool_spans[0] + # Displayed under the AI Config key, not OpenCode's own tool id. + assert span.name == "execute_tool run_shell" + assert span.attributes["gen_ai.tool.name"] == "run_shell" + assert span.attributes["gen_ai.tool.call.id"] == "c1" + assert span.ended == 1 + assert ( + span.status != "ERROR_SENTINEL" + ) # set_status was called with OK, not asserted here + # The NativeTool's own tracking stub fired exactly once. + tool_call_events = [ + call + for call in mock_ld_client.track.call_args_list + if call.args[0] == "$ld:ai:tool_call" + ] + assert len(tool_call_events) == 1 + # `toolKey` is the AI Config key `wrap_tool_handlers` was given, not + # the provider's own tool name. + assert tool_call_events[0].args[2]["toolKey"] == "run_shell" + + async def test_a_failed_tool_call_fails_its_span( + self, monkeypatch: pytest.MonkeyPatch, mock_ld_client: MagicMock + ) -> None: + tracer = self._tracer(monkeypatch) + handler, _script = playing( + [ + assistant_open("m1"), + tool_part("c1", "running", name="bash", raw_input={"command": "rm"}), + tool_part("c1", "error", name="bash", raw_input={"command": "rm"}), + assistant_completed("m1"), + ] + ) + await _drain( + await handler.stream( + _CONFIG_WITH_TOOLS, "hi", _wrapped({"run_shell": OpenCodeBash}), None + ) + ) + tool_spans = [s for s in tracer.spans if s.name.startswith("execute_tool")] + assert len(tool_spans) == 1 + assert len(tool_spans[0].exceptions) == 1 + + async def test_a_chat_span_opens_and_closes_per_assistant_message( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + tracer = self._tracer(monkeypatch) + handler, _script = playing( + [ + assistant_open("m1"), + assistant_completed("m1", tokens={"input": 10, "output": 5}), + assistant_open("m2"), + assistant_completed("m2", tokens={"input": 3, "output": 1}), + ] + ) + await _drain(await handler.stream({}, "hi", None, None)) + chat_spans = [s for s in tracer.spans if s.name.startswith("chat ")] + assert len(chat_spans) == 2 + assert chat_spans[0].attributes["gen_ai.usage.input_tokens"] == 10 + assert chat_spans[0].attributes["gen_ai.usage.output_tokens"] == 5 + assert chat_spans[1].attributes["gen_ai.usage.input_tokens"] == 3 + assert all(span.ended == 1 for span in chat_spans) + + async def test_abandoning_the_stream_marks_open_spans_abandoned_not_failed( + self, monkeypatch: pytest.MonkeyPatch, mock_ld_client: MagicMock + ) -> None: + tracer = self._tracer(monkeypatch) + handler, _script = playing( + [ + assistant_open("m1"), + tool_part("c1", "running", name="bash", raw_input={"command": "sleep"}), + # No completion: the stream is abandoned while this call is open. + ], + ) + agen = ( + await handler.stream( + _CONFIG_WITH_TOOLS, "hi", _wrapped({"run_shell": OpenCodeBash}), None + ) + ).__aiter__() + await agen.__anext__() + await agen.aclose() + tool_spans = [s for s in tracer.spans if s.name.startswith("execute_tool")] + assert len(tool_spans) == 1 + span = tool_spans[0] + assert span.ended == 1 + assert span.attributes.get("launchdarkly.stream.abandoned") is True + assert span.exceptions == [] diff --git a/tests/test_cross_handler_parity.py b/tests/test_cross_handler_parity.py index 7f3e2387..60a3375f 100644 --- a/tests/test_cross_handler_parity.py +++ b/tests/test_cross_handler_parity.py @@ -495,6 +495,11 @@ def test_the_carrier_stays_behind_the_capture_gate(self) -> None: "claude-agents": "launchdarkly_ai_claude_agents.handler", "openai-agents": "launchdarkly_ai_openai_agents.handler", "langchain-agents": "launchdarkly_ai_langchain_agents.handler", + # Not hook-based in the vendor-SDK sense -- there is no vendor SDK here, + # only OpenCode's own bus. The reducer plays the same role: it is the one + # object whose tool-part transitions open and close `execute_tool` spans, + # so it needs the same abandon/close split as a hook object does. + "opencode-agents": "launchdarkly_ai_opencode_agents.handler", }