feat(agent): introduce generalized reverse proxy bidirectional streaming evaluation

evalbench/eval_service.py

@@ -29,12 +29,14 @@
     eval_request_pb2,
     eval_response_pb2,
     eval_service_pb2_grpc,
+    eval_agent_pb2,
 )
 from util.service import (
     load_session_configs,
     get_dataset_from_request,
 )
 from generators.models.grpc_proxy import PROXY_QUEUES
+from generators.models.agentic_reverse_proxy import AGENT_PROXY_QUEUES
 
 import threading
 from util.context import rpc_id_var
@@ -453,6 +455,160 @@ async def read_from_client():
             PROXY_QUEUES.pop(session_id, None)
             logging.info(f"Cleaned up proxy queues for session {session_id}")
 
+    async def AgentInteract(
+        self,
+        request_iterator: AsyncIterator[eval_agent_pb2.AgentStreamMessage],
+        context: grpc.ServicerContext,
+    ) -> AsyncGenerator[eval_agent_pb2.AgentStreamMessage, None]:
+        """Bidirectional stream linking Google3 Autonomous Agents to Evalbench AgentEvaluator."""
+        session_id = rpc_id_var.get()
+        session = SESSIONMANAGER.get_session(session_id)
+        config, db_configs, model_config, setup_config = load_session_configs(session)
+
+        if config is None:
+            context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
+            context.set_details("Session not configured")
+            return
+
+        logging.info("Starting an AgentInteract bidirectional stream for session %s...", session_id)
+        config["session_id"] = session_id
+
+        inboxes: dict[str, queue.Queue] = {}  # correlation_id -> queue.Queue
+        out_queue = queue.Queue()  # Evalbench -> worker
+
+        config["agent_inboxes"] = inboxes
+        config["agent_out_queue"] = out_queue
+        AGENT_PROXY_QUEUES[session_id] = (inboxes, out_queue)
+
+        # Load dataset and instantiate orchestrator
+        dataset_config_json = config.get("dataset_config")
+        dataset_dict = load_dataset_from_json(dataset_config_json, config)
+
+        dataset = []
+        for _, item_list in dataset_dict.items():
+            dataset.extend(item_list)
+
+        num_evals = config.get("num_evals_to_run")
+        if num_evals and int(num_evals) > 0:
+            dataset = dataset[:int(num_evals)]
+
+        orchestrator = get_orchestrator(config, db_configs, setup_config, report_progress=True)
+        loop = asyncio.get_event_loop()
+        ctx = contextvars.copy_context()
+
+        try:
+            def _cleanup_on_drop(ctx):
+                if session_id in AGENT_PROXY_QUEUES:
+                    AGENT_PROXY_QUEUES.pop(session_id, None)
+                    logging.info(f"Cleaned up agent proxy queues for session {session_id}")
+
+            context.add_done_callback(_cleanup_on_drop)
+
+            async def run_eval_and_process():
+                await loop.run_in_executor(None, ctx.run, orchestrator.evaluate, dataset)
+                job_id, run_time, results_tf, scores_tf, multi_trial_scores_tf = orchestrator.process()
+                reporters = get_reporters(config.get("reporting") or {}, job_id, run_time)
+                logging.info("Processing agent evaluation results...")
+                summary = await loop.run_in_executor(
+                    None,
+                    ctx.run,
+                    _process_results,
+                    reporters,
+                    job_id,
+                    run_time,
+                    results_tf,
+                    scores_tf,
+                    multi_trial_scores_tf,
+                    config,
+                    model_config,
+                    db_configs,
+                )
+                return job_id, summary
+
+            eval_task = asyncio.create_task(run_eval_and_process())
+
+            async def read_from_client():
+                async for response in request_iterator:
+                    corr_id = response.correlation_id
+                    logging.info(
+                        "Server-Inbound: Received AgentStreamMessage (correlation_id=%s, payload=%s)",
+                        corr_id,
+                        response.WhichOneof("payload"),
+                    )
+                    if corr_id in inboxes:
+                        inboxes[corr_id].put(response)
+                    else:
+                        logging.warning(
+                            "Server-Inbound: Orphaned AgentStreamMessage correlation_id '%s' (active inboxes: %s)",
+                            corr_id,
+                            list(inboxes.keys()),
+                        )
+
+            read_task = asyncio.create_task(read_from_client())
+
+            # Yield loop: pop from out_queue and yield to worker
+            job_id = None
+            summary = None
+            while True:
+                if eval_task.done():
+                    logging.info("Agent Evaluator & Reporting task finished for session %s.", session_id)
+                    try:
+                        job_id, summary = eval_task.result()
+                    except Exception as e:
+                        logging.error("Agent Evaluator & Reporting task failed: %s", e, exc_info=True)
+                    break
+
+                if SESSIONMANAGER.get_session(session_id) is None:
+                    logging.warning(f"Session {session_id} deleted. Terminating stream.")
+                    context.set_code(grpc.StatusCode.NOT_FOUND)
+                    context.set_details("Session deleted")
+                    return
+
+                try:
+                    out_msg: eval_agent_pb2.AgentStreamMessage = await asyncio.to_thread(out_queue.get, True, 0.5)
+                    logging.info(
+                        "Server-Outbound: Yielding AgentStreamMessage (correlation_id=%s, payload=%s)",
+                        out_msg.correlation_id,
+                        out_msg.WhichOneof("payload"),
+                    )
+                    yield out_msg
+                except queue.Empty:
+                    continue
+                except Exception as e:
+                    logging.error("Server-Outbound: Error yielding message: %s", e, exc_info=True)
+                    continue
+
+            # Flush any remaining messages from out_queue before finishing
+            while not out_queue.empty():
+                try:
+                    out_msg = out_queue.get_nowait()
+                    yield out_msg
+                except Exception:
+                    break
+
+            read_task.cancel()
+            try:
+                await read_task
+            except asyncio.CancelledError:
+                # Expected when cleaning up inbound reader task upon eval completion
+                pass
+
+            if job_id and summary:
+                logging.info(f"Finished Agent Evaluation Job ID {job_id}. Summary: {summary}")
+                final_msg = eval_agent_pb2.AgentStreamMessage(
+                    session_id=session_id,
+                    correlation_id="final_summary",
+                    session_summary=eval_agent_pb2.SessionSummaryMessage(
+                        job_id=job_id,
+                        summary_json=json.dumps(summary),
+                    ),
+                )
+                yield final_msg
+
+        finally:
+            AGENT_PROXY_QUEUES.pop(session_id, None)
+            logging.info(f"Cleaned up agent proxy queues for session {session_id}")
+
 
 def _process_results(
     reporters, job_id, run_time, results_tf, scores_tf, multi_trial_scores_tf, config, model_config, db_configs

evalbench/evalproto/eval_agent.proto

@@ -0,0 +1,142 @@
+edition = "2023";
+
+package cloud_databases_eval_proto;
+
+message AgentStreamMessage {
+  string session_id = 1;
+  string scenario_id = 2;
+  string correlation_id = 3;  // Correlates async requests and responses across threads
+
+  oneof payload {
+    // 0. Pre-Flight Health & Readiness Probe
+    HealthCheckRequest health_check_request = 5;
+    HealthCheckResponse health_check_response = 6;
+
+    // 1. Scenario Lifecycle
+    LifecycleRequest lifecycle_request = 10;
+    LifecycleResponse lifecycle_response = 11;
+
+    // 2. Turn Execution
+    TurnRequest turn_request = 20;
+    TurnResponse turn_response = 21;
+
+    // 3. In-Flight Scoring Execution (Named Remote Scorers)
+    ScoringRequest scoring_request = 30;
+    ScoringResponse scoring_response = 31;
+
+    // 4. Remote Workspace Archival
+    ArtifactRequest artifact_request = 40;
+    ArtifactResponse artifact_response = 41;
+
+    // 5. Final Session Completion
+    SessionSummaryMessage session_summary = 50;
+  }
+}
+
+// --- 0. Health Check Messages ---
+message HealthCheckRequest {
+  string probe_command = 1;
+  float timeout_seconds = 2;
+}
+
+message HealthCheckResponse {
+  bool is_healthy = 1;
+  bool internet_egress_ok = 2;
+  int32 active_skills_count = 3;
+  repeated string available_tools = 4;
+  string diagnostic_output = 5;
+  string error_message = 6;
+}
+
+// --- 1. Lifecycle Messages ---
+message LifecycleRequest {
+  enum Stage { STAGE_UNSPECIFIED = 0; SETUP = 1; RESET = 2; TEARDOWN = 3; }
+  Stage stage = 1;
+  string command = 2;
+  map<string, string> env_vars = 3;
+}
+
+message LifecycleResponse {
+  int32 exit_code = 1;
+  string stdout = 2;
+  string stderr = 3;
+  string error_message = 4;
+}
+
+// --- 2. Turn Execution Messages ---
+message TurnRequest {
+  int32 turn_index = 1;
+  string prompt = 2;
+  string system_instruction = 3;
+  map<string, string> env = 4;
+  string working_dir = 5;
+  float timeout_seconds = 6;
+  bool resume = 7;
+}
+
+message ToolCallRecord {
+  string tool_id = 1;
+  string tool_name = 2;
+  string parameters_json = 3;
+  string output = 4;
+  string status = 5;            // "success", "error"
+  int64 duration_ms = 6;
+}
+
+message TurnResponse {
+  int32 turn_index = 1;
+  string response_text = 2;
+  string stdout = 3;
+  string stderr = 4;
+  int32 exit_code = 5;
+  repeated ToolCallRecord tool_calls = 6;
+  map<string, int64> token_stats = 7;
+  bool execution_completed = 8;
+  string error_message = 9;
+}
+
+// --- 3. In-Flight Scoring Messages ---
+message RemoteScorerSpec {
+  string name = 1;                  // Scorer identifier, e.g. "dataform_compile", "dbt_run", "notebook_eval"
+  string config_json = 2;           // Serialized configuration dictionary from YAML
+  float timeout_seconds = 3;        // Timeout per scorer
+}
+
+message ScoringRequest {
+  repeated RemoteScorerSpec scorers = 1;
+}
+
+message ScoreResult {
+  string name = 1;
+  float score = 2;                  // 100.0 (PASS), 0.0 (FAIL), or continuous score
+  int32 exit_code = 3;
+  string stdout = 4;
+  string stderr = 5;
+  string logs = 6;
+  string error_message = 7;
+}
+
+message ScoringResponse {
+  repeated ScoreResult results = 1;
+}
+
+// --- 4. Artifact Archival Messages ---
+message ArtifactRequest {
+  string target_gcs_bucket = 1;
+  string target_gcs_prefix = 2;
+  string export_path = 3;
+  repeated string exclude_patterns = 4;
+}
+
+message ArtifactResponse {
+  string gcs_uri = 1;
+  int64 archive_size_bytes = 2;
+  repeated string exported_files = 3;
+  string error_message = 4;
+}
+
+// --- 5. Session Summary ---
+message SessionSummaryMessage {
+  string job_id = 1;
+  string summary_json = 2;
+}

evalbench/evalproto/eval_service.proto

@@ -6,6 +6,7 @@ import "eval_config.proto";
 import "eval_connect.proto";
 import "eval_request.proto";
 import "eval_response.proto";
+import "eval_agent.proto";
 
 option java_multiple_files = true;
 
@@ -40,6 +41,11 @@ service EvalService {
     // option deadline = 1800;
   }
 
+  // Autonomous Agent Bidirectional stream.
+  rpc AgentInteract(stream AgentStreamMessage) returns (stream AgentStreamMessage) {
+    // option deadline = 1800;
+  }
+
   // PrepareCodeEvalInputs for NL2Code Evaluation
   rpc PrepareCodeEvalInputs(EvalCodeInputRequest) returns (stream EvalCodeInputRequest) {
     // option deadline = 1800;

evalbench/generators/models/__init__.py

@@ -15,6 +15,7 @@
 from .mcp_tools import McpToolsGenerator
 from .noop_agent import NoopAgentGenerator
 from .agent_runtime import AgentRuntimeGenerator
+from .agentic_reverse_proxy import AgenticReverseProxyGenerator
 from util.config import load_yaml_config
 
 
@@ -34,6 +35,7 @@ def get_generator(global_models, model_config_path: str, db: DB = None):
             "querydata": lambda: QueryData(config),
             "query_data_api": lambda: QueryDataAPIGenerator(config),
             "grpc_proxy": lambda: GrpcProxyModel(config),
+            "agentic_reverse_proxy": lambda: AgenticReverseProxyGenerator(config),
             "gemini_cli": lambda: GeminiCliGenerator(config),
             "claude_code": lambda: ClaudeCodeGenerator(config),
             "codex_cli": lambda: CodexCliGenerator(config),