"""Codex-native goal control helpers for the runner app. The AP server owns the public ``/v1/sessions/{id}/codex_goal`` routes, but the actual goal state lives inside Codex app-server for the loaded Codex thread. This module keeps the runner-side JSON-RPC forwarding and direct ``/events`` validation out of the already-large ``omnigent.runner.app`` module. """ from __future__ import annotations import contextlib import logging from collections.abc import Callable from typing import Any, Protocol from fastapi.responses import JSONResponse, Response from omnigent.harness_plugins import CODEX_NATIVE_CODING_AGENT class BridgeStateForSession(Protocol): """Callable that resolves a live Codex app-server bridge state.""" async def __call__( self, conv_id: str, *, action: str, missing_state_log_level: int = logging.WARNING, ) -> Any | None: ... class ClientSafeErrorDetail(Protocol): """Callable that logs an exception and returns safe client-facing detail.""" def __call__(self, exc: BaseException, *, context: str) -> str: ... class CodexGoalRunner: """Forward Codex goal controls from runner events to Codex app-server.""" def __init__( self, *, bridge_state_for_session: BridgeStateForSession, client_safe_error_detail: ClientSafeErrorDetail, logger: logging.Logger, ) -> None: self._bridge_state_for_session = bridge_state_for_session self._client_safe_error_detail = client_safe_error_detail self._logger = logger @staticmethod def _goal_to_api(goal: dict[str, Any]) -> dict[str, Any]: """ Convert a Codex app-server goal object to Omnigent API field names. Codex app-server speaks camelCase JSON-RPC fields. The AP/UI session API uses snake_case, so normalize the goal at the runner boundary before returning it through ``/events``. """ def required_str(name: str) -> str: value = goal.get(name) if not isinstance(value, str): raise ValueError(f"malformed upstream goal: missing string {name!r}") return value def required_non_negative_int(name: str) -> int: value = goal.get(name) if isinstance(value, bool) or not isinstance(value, int) or value < 0: raise ValueError(f"malformed upstream goal: missing non-negative integer {name!r}") return value def optional_int(name: str) -> int | None: value = goal.get(name) if value is None: return None if isinstance(value, bool) or not isinstance(value, int): raise ValueError(f"malformed upstream goal: invalid integer {name!r}") return value token_budget = goal.get("tokenBudget") if token_budget is not None and ( isinstance(token_budget, bool) or not isinstance(token_budget, int) or token_budget <= 0 ): raise ValueError("malformed upstream goal: invalid integer 'tokenBudget'") return { "thread_id": required_str("threadId"), "objective": required_str("objective"), "status": required_str("status"), "token_budget": token_budget, "tokens_used": required_non_negative_int("tokensUsed"), "time_used_seconds": required_non_negative_int("timeUsedSeconds"), "created_at": optional_int("createdAt"), "updated_at": optional_int("updatedAt"), } @staticmethod def _goal_result(response: dict[str, Any], *, action: str) -> dict[str, Any]: """Extract a JSON-RPC ``result`` object from a Codex goal response.""" result = response.get("result") if not isinstance(result, dict): raise ValueError(f"Codex goal {action} returned no result object") return result @staticmethod def _bridge_error(action: str) -> JSONResponse: """Build the standard no-bridge response for Codex goal operations.""" return JSONResponse( status_code=502, content={ "error": "codex_native_goal_failed", "detail": (f"Codex-native goal {action} requires a loaded Codex bridge."), }, ) @staticmethod def _malformed_response(action: str, detail: str) -> JSONResponse: """Build the standard malformed-response error for Codex goal operations.""" return JSONResponse( status_code=503, content={ "error": "codex_native_goal_failed", "detail": f"Codex-native goal {action} returned {detail}.", }, ) async def _request( self, conv_id: str, *, action: str, method: str, params: dict[str, Any], ) -> dict[str, Any] | JSONResponse: """Execute one Codex app-server goal JSON-RPC request.""" from omnigent.codex_native_app_server import client_for_transport state = await self._bridge_state_for_session(conv_id, action=f"goal {action}") if state is None: return self._bridge_error(action) codex_client = client_for_transport( state.socket_path, client_name="omnigent-codex-native-runner", ) try: await codex_client.connect() return self._goal_result( await codex_client.request(method, {"threadId": state.thread_id, **params}), action=action, ) except ValueError as exc: return self._malformed_response(action, str(exc)) except Exception as exc: # noqa: BLE001 - surface app-server goal failures to AP. self._logger.warning( "Codex-native %s failed for session=%s", method, conv_id, exc_info=True, ) return JSONResponse( status_code=503, content={ "error": "codex_native_goal_failed", "detail": self._client_safe_error_detail( exc, context=f"codex-native goal {action}", ), }, ) finally: with contextlib.suppress(Exception): await codex_client.close() async def get(self, conv_id: str) -> Response: """Read the current Codex app-server goal for a loaded thread.""" result = await self._request( conv_id, action="read", method="thread/goal/get", params={}, ) if isinstance(result, JSONResponse): return result goal = result.get("goal") if goal is not None and not isinstance(goal, dict): return self._malformed_response("read", "invalid goal object") try: api_goal = None if goal is None else self._goal_to_api(goal) except ValueError as exc: return self._malformed_response("read", str(exc)) return JSONResponse({"goal": api_goal}) async def set( self, conv_id: str, *, objective: str, token_budget: int | None, token_budget_provided: bool, status: str | None, ) -> Response: """Create or replace the current Codex app-server goal for a thread.""" params: dict[str, Any] = { "objective": objective, } if token_budget_provided: params["tokenBudget"] = token_budget if status is not None: params["status"] = status result = await self._request( conv_id, action="set", method="thread/goal/set", params=params, ) if isinstance(result, JSONResponse): return result goal = result.get("goal") if not isinstance(goal, dict): return self._malformed_response("set", "invalid goal object") try: api_goal = self._goal_to_api(goal) except ValueError as exc: return self._malformed_response("set", str(exc)) return JSONResponse({"goal": api_goal}) async def update_status(self, conv_id: str, *, status: str) -> Response: """ Update the current Codex app-server goal status. Codex models pause/resume as ``thread/goal/set`` calls with only the status field. Objective and token-budget fields are omitted so Codex keeps the existing goal text and budget. """ result = await self._request( conv_id, action="update status", method="thread/goal/set", params={"status": status}, ) if isinstance(result, JSONResponse): return result goal = result.get("goal") if not isinstance(goal, dict): return self._malformed_response( "update status", "invalid goal object", ) try: api_goal = self._goal_to_api(goal) except ValueError as exc: return self._malformed_response("update status", str(exc)) return JSONResponse({"goal": api_goal}) async def clear(self, conv_id: str) -> Response: """Clear the current Codex app-server goal for a loaded thread.""" result = await self._request( conv_id, action="clear", method="thread/goal/clear", params={}, ) if isinstance(result, JSONResponse): return result cleared = result.get("cleared") if not isinstance(cleared, bool): return self._malformed_response("clear", "invalid cleared flag") return JSONResponse({"cleared": cleared}) async def handle_event( self, conv_id: str, body_type: str | None, body: Any, *, session_harness_name: Callable[[str], str | None], ) -> Response | None: """Handle direct runner ``/events`` payloads for Codex goal controls.""" if body_type not in {"goal_get", "goal_set", "goal_status", "goal_clear"}: return None if session_harness_name(conv_id) != CODEX_NATIVE_CODING_AGENT.harness: return JSONResponse( status_code=400, content={ "error": "invalid_input", "detail": "Codex goal controls require a codex-native session", }, ) if body_type == "goal_get": return await self.get(conv_id) if body_type == "goal_clear": return await self.clear(conv_id) if body_type == "goal_set": # The AP route validates this too, but direct runner callers should # get a clear 400 instead of an app-server validation error. objective = body.get("objective") if isinstance(body, dict) else None if not isinstance(objective, str) or not objective.strip(): return JSONResponse( status_code=400, content={ "error": "invalid_input", "detail": "Body 'objective' must be a non-empty string", }, ) if len(objective) > 4000: return JSONResponse( status_code=400, content={ "error": "invalid_input", "detail": "Body 'objective' must be at most 4000 characters", }, ) token_budget_provided = isinstance(body, dict) and "token_budget" in body token_budget = body.get("token_budget") if token_budget_provided else None if token_budget is not None and ( isinstance(token_budget, bool) or not isinstance(token_budget, int) or token_budget <= 0 ): return JSONResponse( status_code=400, content={ "error": "invalid_input", "detail": "Body 'token_budget' must be a positive integer or null", }, ) status = body.get("status") if isinstance(body, dict) else None if status is not None and status not in ("active", "paused"): return JSONResponse( status_code=400, content={ "error": "invalid_input", "detail": "Body 'status' must be 'active' or 'paused'", }, ) return await self.set( conv_id, objective=objective.strip(), token_budget=token_budget, token_budget_provided=token_budget_provided, status=status, ) status = body.get("status") if isinstance(body, dict) else None if status not in ("active", "paused"): return JSONResponse( status_code=400, content={ "error": "invalid_input", "detail": "Body 'status' must be 'active' or 'paused'", }, ) return await self.update_status(conv_id, status=status)