nighthawk-python 0.1.0__py3-none-any.whl

nighthawk/backends/mcp_boundary.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import json
+
+import tiktoken
+from opentelemetry import context as otel_context
+from opentelemetry.context import Context as OtelContext
+
+from ..runtime.step_context import DEFAULT_TOOL_RESULT_RENDERING_POLICY, ToolResultRenderingPolicy, get_current_step_context
+from ..tools.contracts import render_tool_result_json_text
+from .tool_bridge import ToolHandler
+
+
+def _resolve_tool_result_rendering_policy() -> ToolResultRenderingPolicy:
+    try:
+        step_context = get_current_step_context()
+    except Exception:
+        return DEFAULT_TOOL_RESULT_RENDERING_POLICY
+
+    if step_context.tool_result_rendering_policy is None:
+        return DEFAULT_TOOL_RESULT_RENDERING_POLICY
+    return step_context.tool_result_rendering_policy
+
+
+def _tool_boundary_failure_text(*, message: str, guidance: str) -> str:
+    rendering_policy = _resolve_tool_result_rendering_policy()
+    encoding = tiktoken.get_encoding(rendering_policy.tokenizer_encoding_name)
+    return render_tool_result_json_text(
+        value=None,
+        error={"kind": "internal", "message": message, "guidance": guidance},
+        max_tokens=rendering_policy.tool_result_max_tokens,
+        encoding=encoding,
+        style=rendering_policy.json_renderer_style,
+    )
+
+
+def _minimal_tool_boundary_failure_json_text(*, message: str, guidance: str) -> str:
+    payload = {
+        "value": None,
+        "error": {
+            "kind": "internal",
+            "message": message,
+            "guidance": guidance,
+        },
+    }
+    return json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
+
+
+async def _call_tool_handler_result_text(
+    *,
+    tool_name: str,
+    arguments: dict[str, object],
+    tool_handler: ToolHandler,
+    parent_otel_context: OtelContext,
+) -> str:
+    _ = tool_name
+
+    context_token = otel_context.attach(parent_otel_context)
+    try:
+        return await tool_handler(arguments)
+    except Exception as exception:
+        error_message = str(exception) or "Tool boundary wrapper failed"
+        try:
+            return _tool_boundary_failure_text(
+                message=error_message,
+                guidance="The tool boundary wrapper failed. Retry or report this error.",
+            )
+        except Exception:
+            return _minimal_tool_boundary_failure_json_text(
+                message=error_message,
+                guidance="The tool boundary wrapper failed and could not access the environment. Retry or report this error.",
+            )
+    finally:
+        otel_context.detach(context_token)
+
+
+async def _get_safe_tool_result_text(
+    *,
+    tool_name: str,
+    arguments: dict[str, object],
+    tool_handler: ToolHandler,
+    parent_otel_context: OtelContext,
+) -> str:
+    try:
+        return await _call_tool_handler_result_text(
+            tool_name=tool_name,
+            arguments=arguments,
+            tool_handler=tool_handler,
+            parent_otel_context=parent_otel_context,
+        )
+    except Exception as exception:
+        return _minimal_tool_boundary_failure_json_text(
+            message=str(exception) or "Tool boundary wrapper failed",
+            guidance="The tool boundary wrapper failed. Retry or report this error.",
+        )
+
+
+async def call_tool_for_claude_code_sdk(
+    *,
+    tool_name: str,
+    arguments: dict[str, object],
+    tool_handler: ToolHandler,
+    parent_otel_context: OtelContext,
+) -> dict[str, object]:
+    result_text = await _get_safe_tool_result_text(
+        tool_name=tool_name,
+        arguments=arguments,
+        tool_handler=tool_handler,
+        parent_otel_context=parent_otel_context,
+    )
+    return {"content": [{"type": "text", "text": result_text}]}
+
+
+async def call_tool_for_low_level_mcp_server(
+    *,
+    tool_name: str,
+    arguments: dict[str, object],
+    tool_handler: ToolHandler,
+    parent_otel_context: OtelContext,
+) -> list[object]:
+    from mcp import types as mcp_types
+
+    result_text = await _get_safe_tool_result_text(
+        tool_name=tool_name,
+        arguments=arguments,
+        tool_handler=tool_handler,
+        parent_otel_context=parent_otel_context,
+    )
+    return [mcp_types.TextContent(type="text", text=result_text)]

nighthawk/backends/mcp_server.py

@@ -0,0 +1,226 @@
+"""Embedded MCP tool server for provider backends.
+
+Starts a Streamable HTTP MCP server in a background thread, exposing
+Nighthawk tools to CLI-based backends (e.g. Codex) that consume tools
+via an MCP endpoint.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import socket
+import threading
+import time
+from collections.abc import AsyncIterator, Awaitable, Callable
+from contextlib import asynccontextmanager
+from contextvars import copy_context
+from typing import Any, cast
+
+import tiktoken
+from opentelemetry import context as otel_context
+from pydantic_ai import RunContext
+from pydantic_ai.exceptions import UnexpectedModelBehavior
+from pydantic_ai.tools import ToolDefinition
+
+from ..runtime.step_context import (
+    DEFAULT_TOOL_RESULT_RENDERING_POLICY,
+    StepContext,
+    ToolResultRenderingPolicy,
+)
+from .mcp_boundary import call_tool_for_low_level_mcp_server
+
+
+class McpServer:
+    """In-process MCP tool server backed by Nighthawk tool handlers."""
+
+    def __init__(
+        self,
+        *,
+        tool_name_to_tool_definition: dict[str, ToolDefinition],
+        tool_name_to_handler: dict[str, Callable[[dict[str, Any]], Awaitable[str]]],
+        tool_result_rendering_policy: ToolResultRenderingPolicy,
+        parent_otel_context: Any,
+    ) -> None:
+        self._tool_name_to_tool_definition = tool_name_to_tool_definition
+        self._tool_name_to_handler = tool_name_to_handler
+        self._tool_result_rendering_policy = tool_result_rendering_policy
+        self._parent_otel_context = parent_otel_context
+
+        self._server: Any | None = None
+        self._server_thread: threading.Thread | None = None
+        self._listening_socket: socket.socket | None = None
+        self._url: str | None = None
+
+    @property
+    def url(self) -> str:
+        if self._url is None:
+            raise RuntimeError("MCP tool server is not started")
+        return self._url
+
+    def start(self) -> None:
+        if self._server_thread is not None:
+            raise RuntimeError("MCP tool server is already started")
+
+        import uvicorn
+        from mcp.server.fastmcp.server import StreamableHTTPASGIApp
+        from mcp.server.lowlevel.server import Server as McpLowLevelServer
+        from mcp.server.streamable_http_manager import (
+            StreamableHTTPSessionManager,
+        )
+        from starlette.applications import Starlette
+        from starlette.routing import Route
+
+        mcp_server = McpLowLevelServer("nighthawk")
+
+        @mcp_server.list_tools()
+        async def list_tools() -> list[Any]:
+            from mcp import types as mcp_types
+
+            tools: list[mcp_types.Tool] = []
+            for tool_name in sorted(self._tool_name_to_handler.keys()):
+                tool_definition = self._tool_name_to_tool_definition.get(tool_name)
+                if tool_definition is None:
+                    raise RuntimeError(f"Tool definition missing for {tool_name!r}")
+
+                tools.append(
+                    mcp_types.Tool(
+                        name=tool_name,
+                        description=tool_definition.description or "",
+                        inputSchema=tool_definition.parameters_json_schema,
+                    )
+                )
+            return tools
+
+        @mcp_server.call_tool(validate_input=False)
+        async def call_tool(name: str, arguments: dict[str, Any]) -> Any:
+            from mcp import types as mcp_types
+
+            from ..tools.contracts import render_tool_result_json_text
+
+            handler = self._tool_name_to_handler.get(name)
+            if handler is None:
+                policy = self._tool_result_rendering_policy
+                result_text = render_tool_result_json_text(
+                    value=None,
+                    error={"kind": "resolution", "message": f"Unknown tool: {name}", "guidance": "Choose a visible tool name and retry."},
+                    max_tokens=policy.tool_result_max_tokens,
+                    encoding=tiktoken.get_encoding(policy.tokenizer_encoding_name),
+                    style=policy.json_renderer_style,
+                )
+                return [mcp_types.TextContent(type="text", text=result_text)]
+
+            return await call_tool_for_low_level_mcp_server(
+                tool_name=name,
+                arguments=arguments,
+                tool_handler=handler,
+                parent_otel_context=self._parent_otel_context,
+            )
+
+        session_manager = StreamableHTTPSessionManager(app=mcp_server)
+        streamable_http_asgi = StreamableHTTPASGIApp(session_manager)
+
+        starlette_application = Starlette(
+            routes=[Route("/mcp", endpoint=streamable_http_asgi)],
+            lifespan=lambda _app: session_manager.run(),
+        )
+
+        listening_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        listening_socket.bind(("127.0.0.1", 0))
+        listening_socket.listen(128)
+        host, port = listening_socket.getsockname()
+
+        configuration = uvicorn.Config(
+            starlette_application,
+            host=str(host),
+            port=int(port),
+            log_level="warning",
+            lifespan="on",
+        )
+        server = uvicorn.Server(configuration)
+
+        def run_server() -> None:
+            server.run(sockets=[listening_socket])
+
+        context = copy_context()
+        server_thread = threading.Thread(
+            target=context.run,
+            args=(run_server,),
+            name="nighthawk-mcp-server",
+            daemon=True,
+        )
+        server_thread.start()
+
+        self._server = server
+        self._server_thread = server_thread
+        self._listening_socket = listening_socket
+        self._url = f"http://127.0.0.1:{port}/mcp"
+
+    async def stop(self) -> None:
+        if self._server is None or self._server_thread is None:
+            return
+
+        self._server.should_exit = True
+        await asyncio.to_thread(self._server_thread.join, 5)
+
+        if self._listening_socket is not None:
+            with contextlib.suppress(Exception):
+                self._listening_socket.close()
+
+        self._server = None
+        self._server_thread = None
+        self._listening_socket = None
+        self._url = None
+
+
+async def _wait_for_tcp_listen(host: str, port: int, *, timeout_seconds: float) -> None:
+    deadline = time.monotonic() + timeout_seconds
+    while time.monotonic() < deadline:
+        try:
+            reader, writer = await asyncio.open_connection(host, port)
+        except OSError:
+            await asyncio.sleep(0.01)
+            continue
+        else:
+            writer.close()
+            await writer.wait_closed()
+            return
+    raise UnexpectedModelBehavior("Timed out waiting for MCP tool server to start")
+
+
+@asynccontextmanager
+async def mcp_server_if_needed(
+    *,
+    tool_name_to_tool_definition: dict[str, ToolDefinition],
+    tool_name_to_handler: dict[str, Callable[[dict[str, Any]], Awaitable[str]]],
+) -> AsyncIterator[str | None]:
+    if not tool_name_to_handler:
+        yield None
+        return
+
+    from pydantic_ai._run_context import get_current_run_context
+
+    parent_otel_context = otel_context.get_current()
+    parent_run_context = get_current_run_context()
+    if parent_run_context is None:
+        raise RuntimeError("Codex MCP tool server requires an active RunContext")
+    typed_parent_run_context = cast(RunContext[StepContext], parent_run_context)
+    tool_result_rendering_policy = typed_parent_run_context.deps.tool_result_rendering_policy
+    if tool_result_rendering_policy is None:
+        tool_result_rendering_policy = DEFAULT_TOOL_RESULT_RENDERING_POLICY
+
+    server = McpServer(
+        tool_name_to_tool_definition=tool_name_to_tool_definition,
+        tool_name_to_handler=tool_name_to_handler,
+        tool_result_rendering_policy=tool_result_rendering_policy,
+        parent_otel_context=parent_otel_context,
+    )
+    server.start()
+
+    try:
+        url = server.url
+        port = int(url.split(":")[2].split("/", 1)[0])
+        await _wait_for_tcp_listen("127.0.0.1", port, timeout_seconds=2.0)
+        yield url
+    finally:
+        await server.stop()

nighthawk/backends/tool_bridge.py

@@ -0,0 +1,240 @@
+"""Bridge between Nighthawk tools and provider backends.
+
+Provides tool handler construction and tool execution for backends that
+expose Nighthawk tools through Pydantic AI FunctionToolset.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Any, cast
+
+import tiktoken
+from pydantic_ai import RunContext
+
+# NOTE: pydantic_ai private API — no public alternative for custom backends
+# that need to set RunContext for tool execution. Monitor pydantic_ai releases.
+from pydantic_ai._run_context import set_current_run_context
+from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UnexpectedModelBehavior, UserError
+from pydantic_ai.messages import RetryPromptPart
+from pydantic_ai.models import ModelRequestParameters
+from pydantic_ai.toolsets.function import FunctionToolset
+
+from ..runtime.step_context import DEFAULT_TOOL_RESULT_RENDERING_POLICY, StepContext, ToolResultRenderingPolicy
+from ..runtime.tool_calls import generate_tool_call_id, run_tool_instrumented
+from ..tools.contracts import (
+    ToolBoundaryError,
+    ToolResult,
+    render_tool_result_json_text,
+)
+from ..tools.execution import ToolResultWrapperToolset
+
+type ToolHandler = Callable[[dict[str, Any]], Awaitable[str]]
+
+
+def resolve_tool_result_rendering_policy(run_context: RunContext[StepContext]) -> ToolResultRenderingPolicy:
+    policy = run_context.deps.tool_result_rendering_policy
+    if policy is None:
+        return DEFAULT_TOOL_RESULT_RENDERING_POLICY
+    return policy
+
+
+def get_current_run_context_required() -> RunContext[StepContext]:
+    from pydantic_ai._run_context import get_current_run_context
+
+    run_context = get_current_run_context()
+    if run_context is None:
+        raise RuntimeError("Nighthawk tool boundaries require an active RunContext")
+    return cast(RunContext[StepContext], run_context)
+
+
+def resolve_allowed_tool_names(
+    *,
+    model_request_parameters: ModelRequestParameters,
+    configured_allowed_tool_names: tuple[str, ...] | None,
+    available_tool_names: tuple[str, ...],
+) -> tuple[str, ...]:
+    if configured_allowed_tool_names is None:
+        return tuple(
+            name for name in (tool_definition.name for tool_definition in model_request_parameters.function_tools) if name in available_tool_names
+        )
+
+    unknown = [name for name in configured_allowed_tool_names if name not in available_tool_names]
+    if unknown:
+        unknown_list = ", ".join(repr(name) for name in unknown)
+        raise ValueError(f"Configured allowed_tool_names includes unknown tools: {unknown_list}")
+
+    return configured_allowed_tool_names
+
+
+async def execute_tool_call(
+    *,
+    tool_name: str,
+    tool: Any,
+    arguments: dict[str, Any],
+    run_context: RunContext[StepContext],
+    tool_call_id: str,
+    rendering_policy: ToolResultRenderingPolicy,
+    encoding: tiktoken.Encoding,
+) -> str:
+    """Execute a single tool call: validate arguments, call the toolset, and render the result."""
+    with set_current_run_context(run_context):
+        try:
+            validated_arguments = tool.args_validator.validate_python(arguments)
+        except Exception as exception:
+            errors_method = getattr(exception, "errors", None)
+            if callable(errors_method):
+                try:
+                    error_details = errors_method(include_url=False, include_context=False)
+                    if isinstance(error_details, list):
+                        return RetryPromptPart(
+                            tool_name=tool_name,
+                            content=error_details,
+                            tool_call_id=tool_call_id,
+                        ).model_response()
+                except Exception:
+                    pass
+
+            return RetryPromptPart(
+                tool_name=tool_name,
+                content=str(exception) or "Invalid tool arguments",
+                tool_call_id=tool_call_id,
+            ).model_response()
+
+        try:
+            tool_result = await tool.toolset.call_tool(tool_name, validated_arguments, run_context, tool)
+        except (ModelRetry, CallDeferred, ApprovalRequired):
+            raise
+        except ToolBoundaryError as exception:
+            return render_tool_result_json_text(
+                value=None,
+                error={"kind": exception.kind, "message": str(exception), "guidance": exception.guidance},
+                max_tokens=rendering_policy.tool_result_max_tokens,
+                encoding=encoding,
+                style=rendering_policy.json_renderer_style,
+            )
+        except (UserError, UnexpectedModelBehavior) as exception:
+            return render_tool_result_json_text(
+                value=None,
+                error={"kind": "internal", "message": str(exception), "guidance": "The tool backend failed. Retry or report this error."},
+                max_tokens=rendering_policy.tool_result_max_tokens,
+                encoding=encoding,
+                style=rendering_policy.json_renderer_style,
+            )
+        except Exception as exception:
+            return render_tool_result_json_text(
+                value=None,
+                error={
+                    "kind": "internal",
+                    "message": str(exception) or "Tool execution failed",
+                    "guidance": "The tool execution raised an unexpected error. Retry or report this error.",
+                },
+                max_tokens=rendering_policy.tool_result_max_tokens,
+                encoding=encoding,
+                style=rendering_policy.json_renderer_style,
+            )
+
+    if isinstance(tool_result, ToolResult):
+        return render_tool_result_json_text(
+            value=tool_result.value,
+            error=None,
+            max_tokens=rendering_policy.tool_result_max_tokens,
+            encoding=encoding,
+            style=rendering_policy.json_renderer_style,
+        )
+
+    return render_tool_result_json_text(
+        value=tool_result,
+        error=None,
+        max_tokens=rendering_policy.tool_result_max_tokens,
+        encoding=encoding,
+        style=rendering_policy.json_renderer_style,
+    )
+
+
+async def build_tool_name_to_handler(
+    *,
+    model_request_parameters: ModelRequestParameters,
+    visible_tools: list[Any],
+) -> dict[str, ToolHandler]:
+    run_context = get_current_run_context_required()
+
+    toolset = ToolResultWrapperToolset(FunctionToolset(visible_tools))
+
+    tool_name_to_tool = await toolset.get_tools(run_context)
+
+    function_tool_names = {tool_definition.name for tool_definition in model_request_parameters.function_tools}
+
+    tool_name_to_handler: dict[str, ToolHandler] = {}
+
+    for tool_name, tool in tool_name_to_tool.items():
+        if tool_name not in function_tool_names:
+            continue
+
+        async def handler(
+            arguments: dict[str, Any],
+            *,
+            tool_name: str = tool_name,
+            tool: Any = tool,
+        ) -> str:
+            tool_call_id = generate_tool_call_id()
+            from dataclasses import replace
+
+            tool_run_context = replace(
+                run_context,
+                tool_name=tool_name,
+                tool_call_id=tool_call_id,
+            )
+
+            rendering_policy = resolve_tool_result_rendering_policy(run_context)
+            encoding = tiktoken.get_encoding(rendering_policy.tokenizer_encoding_name)
+
+            async def call() -> str:
+                return await execute_tool_call(
+                    tool_name=tool_name,
+                    tool=tool,
+                    arguments=arguments,
+                    run_context=tool_run_context,
+                    tool_call_id=tool_call_id,
+                    rendering_policy=rendering_policy,
+                    encoding=encoding,
+                )
+
+            return await run_tool_instrumented(
+                tool_name=tool_name,
+                arguments=arguments,
+                call=call,
+                run_context=tool_run_context,
+                tool_call_id=tool_call_id,
+            )
+
+        tool_name_to_handler[tool_name] = handler
+
+    return tool_name_to_handler
+
+
+async def prepare_allowed_tools(
+    *,
+    model_request_parameters: ModelRequestParameters,
+    configured_allowed_tool_names: tuple[str, ...] | None,
+    visible_tools: list[Any],
+) -> tuple[dict[str, Any], dict[str, ToolHandler], tuple[str, ...]]:
+    """Build tool definitions, handlers, and allowed names for a backend request."""
+    tool_name_to_handler = await build_tool_name_to_handler(
+        model_request_parameters=model_request_parameters,
+        visible_tools=visible_tools,
+    )
+    available_tool_names = tuple(tool_name_to_handler.keys())
+
+    allowed_tool_names = resolve_allowed_tool_names(
+        model_request_parameters=model_request_parameters,
+        configured_allowed_tool_names=configured_allowed_tool_names,
+        available_tool_names=available_tool_names,
+    )
+
+    tool_name_to_tool_definition = {
+        name: tool_definition for name, tool_definition in model_request_parameters.tool_defs.items() if name in allowed_tool_names
+    }
+    tool_name_to_handler = {name: tool_name_to_handler[name] for name in allowed_tool_names}
+
+    return tool_name_to_tool_definition, tool_name_to_handler, allowed_tool_names