File size: 4,504 Bytes
7698d12
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
"""OpenCode Environment Client.

Thin MCP client over the deployed ``opencode-openenv`` server. The server
exposes a single tool ``run_rollout`` that takes a task + LLM endpoint,
runs one OpenCode agent rollout in a fresh E2B sandbox, and returns a
JSON-serialized :class:`RolloutResult`.

Example::

    from opencode_env_server import OpenCodeEnv

    with OpenCodeEnv(base_url="https://adithya-s-k-opencode-openenv.hf.space") as env:
        env.reset()
        result = env.run_rollout(
            vllm_url="https://your-llm-host/v1",
            model="Qwen/Qwen3.5-4B",
            instruction="Write fizzbuzz.py in the current directory.",
            test_script="#!/bin/bash\\n...",
            task_id="fizzbuzz_001",
            mode="transparent_proxy",
            disable_thinking=True,
        )
        print(result.reward, len(result.proxy_turns))

Docker convenience::

    env = OpenCodeEnv.from_docker_image("opencode-openenv:latest")
    env.reset()
    ...
"""

from __future__ import annotations

import json
from typing import Any

from openenv.core.mcp_client import MCPToolClient

try:
    from .models import RolloutResult
except ImportError:
    from models import RolloutResult  # type: ignore


class OpenCodeEnv(MCPToolClient):
    """Client for the OpenCode OpenEnv server.

    Inherits MCP plumbing (``reset``, ``call_tool``, ``list_tools``,
    ``from_docker_image``, context-manager support) from
    :class:`MCPToolClient`. Adds :meth:`run_rollout` as a typed helper that
    deserializes the tool result into a :class:`RolloutResult`.
    """

    def run_rollout(
        self,
        *,
        vllm_url: str,
        model: str,
        instruction: str,
        test_script: str,
        task_id: str = "",
        setup_shell: str = "",
        upload_files: dict[str, str] | None = None,
        provider: str = "openai_compatible",
        api_key: str = "intercepted",
        mode: str = "transparent_proxy",
        disable_thinking: bool = False,
        max_tokens_cap: int = 4096,
        agent_timeout_s: float = 600.0,
    ) -> RolloutResult:
        """Typed helper around ``call_tool("run_rollout", ...)``."""

        raw = self.call_tool(
            "run_rollout",
            vllm_url=vllm_url,
            model=model,
            instruction=instruction,
            test_script=test_script,
            task_id=task_id,
            setup_shell=setup_shell,
            upload_files=upload_files or {},
            provider=provider,
            api_key=api_key,
            mode=mode,
            disable_thinking=disable_thinking,
            max_tokens_cap=max_tokens_cap,
            agent_timeout_s=agent_timeout_s,
        )
        payload = _extract_text(raw)
        return RolloutResult.model_validate_json(payload)


def _extract_text(result: Any) -> str:
    """Pull the text payload out of an MCP tool result shape.

    Handles three shapes MCPToolClient / call_tool may return:
    - the raw tool text (str)
    - a CallToolObservation-like object with ``.result.content[0].text``
    - a dict with ``content`` list containing ``{"text": ...}`` entries
    """

    if isinstance(result, str):
        return result

    # Object with attribute chain: obs.result.content[0].text
    inner = getattr(result, "result", None)
    if inner is not None:
        content = getattr(inner, "content", None)
        if content:
            first = content[0]
            text = getattr(first, "text", None)
            if isinstance(text, str):
                return text
            if isinstance(first, dict) and "text" in first:
                return first["text"]

    if isinstance(result, dict):
        content = result.get("content")
        if isinstance(content, list) and content:
            first = content[0]
            if isinstance(first, dict) and "text" in first:
                return first["text"]
        nested = result.get("result")
        if isinstance(nested, dict):
            content = nested.get("content")
            if isinstance(content, list) and content:
                first = content[0]
                if isinstance(first, dict) and "text" in first:
                    return first["text"]
        return json.dumps(result, default=str)

    # Object with .content directly
    content = getattr(result, "content", None)
    if content:
        first = content[0]
        text = getattr(first, "text", None)
        if isinstance(text, str):
            return text

    return str(result)