diff --git a/envs/opencode_env/harness.py b/envs/opencode_env/harness.py index a87017e9c..4bfecb0cc 100644 --- a/envs/opencode_env/harness.py +++ b/envs/opencode_env/harness.py @@ -189,7 +189,7 @@ def fetch_proxy_trace(self) -> list[dict[str, Any]]: return records -class OpenCodeSessionFactory(ResourceSessionFactory): +class OpenCodeSessionFactory(ResourceSessionFactory[OpenCodeSession]): """Produce isolated per-rollout :class:`OpenCodeSession` instances. The factory owns sandbox provisioning, opencode install, config injection, diff --git a/src/openenv/core/harness/__init__.py b/src/openenv/core/harness/__init__.py index ece1d3936..34d42f692 100644 --- a/src/openenv/core/harness/__init__.py +++ b/src/openenv/core/harness/__init__.py @@ -12,7 +12,7 @@ import math from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Any, Callable, Protocol +from typing import Any, Callable, Generic, Protocol, TypeVar from ..client_types import StepResult from ..env_server.mcp_types import JsonRpcErrorCode, JsonRpcResponse, Tool @@ -141,8 +141,17 @@ def close(self) -> None: """Release session resources.""" -class ResourceSessionFactory(ABC): - """Factory for producing isolated per-rollout sessions.""" +SessionT = TypeVar("SessionT", bound=ResourceSession) + + +class ResourceSessionFactory(ABC, Generic[SessionT]): + """Factory for producing isolated per-rollout sessions. + + Generic over the concrete session type it creates, so a subclass such as + ``OpenCodeSessionFactory(ResourceSessionFactory[OpenCodeSession])`` types + ``create`` as returning ``OpenCodeSession``. Subclassing without a type + argument stays valid and behaves as before. + """ @abstractmethod def create( @@ -150,7 +159,7 @@ def create( task: Any, seed: int | None = None, episode_id: str | None = None, - ) -> ResourceSession: + ) -> SessionT: """Create one isolated resource session for a rollout."""