33from __future__ import annotations
44
55import json
6+ import os
67import shlex
8+ import time
79from typing import Any , override
810
911from harbor .agents .base import BaseAgent
@@ -24,6 +26,28 @@ def _is_reasoning_model(model: str) -> bool:
2426MAX_TURNS = 32
2527MAX_TOOL_OUTPUT_CHARS = 40_000
2628
29+ # Wall clock, not turn count, is this agent's binding constraint. Most of a turn
30+ # is the sequential `search.py` invocations, each of which boots a fresh
31+ # pyserini/Lucene JVM against the BM25 index (~13s apiece), against ~3-9s of
32+ # inference. Trials therefore run out of clock long before MAX_TURNS, so the
33+ # turn-budget fallback below never fires and harbor instead kills the trial
34+ # mid-search with no answer recorded. This gives that same best-effort safety
35+ # net a deadline to watch. Harbor does not tell an agent its own timeout, so the
36+ # budget is read from the environment; keep it in step with whatever
37+ # --agent-timeout-multiplier the run uses.
38+ AGENT_BUDGET_SEC = float (os .environ .get ("BROWSECOMP_AGENT_BUDGET_SEC" , "900" ))
39+ # Enough to prefill the accumulated history and write a short answer (~8s for a
40+ # 120k-token prompt measured against the target endpoint), plus the answer
41+ # upload and margin for one retry.
42+ FINAL_RESERVE_SEC = float (os .environ .get ("BROWSECOMP_FINAL_RESERVE_SEC" , "120" ))
43+ # A malformed, empty, or length-truncated assistant turn is recoverable; nudge
44+ # rather than abort, and only give up after this many consecutive nudges.
45+ MAX_EMPTY_RETRIES = 3
46+ # The required answer is three short labeled lines, so the forced-final call is
47+ # capped well below the per-turn budget to stop the model from spending its last
48+ # seconds on chain-of-thought and getting truncated a second time.
49+ ANSWER_MAX_TOKENS = 3_000
50+
2751INSTRUCTIONS = """You are a persistent deep-research agent working only against the
2852fixed BrowseComp-Plus corpus. Find the precise answer by issuing focused searches,
2953opening promising documents, reformulating queries, and cross-checking evidence.
@@ -133,13 +157,17 @@ def _trace(self, event: dict[str, Any]) -> None:
133157 file .write (json .dumps (event , ensure_ascii = False , default = str ) + "\n " )
134158
135159 async def _index_command (
136- self , environment : BaseEnvironment , command : str , value : str
160+ self ,
161+ environment : BaseEnvironment ,
162+ command : str ,
163+ value : str ,
164+ timeout_sec : float = 120 ,
137165 ) -> dict [str , Any ]:
138166 flag = "--query" if command == "search" else "--docid"
139167 result = await environment .exec (
140168 f"python3 /opt/browsecomp/search.py { command } { flag } { shlex .quote (value )} " ,
141169 cwd = "/app" ,
142- timeout_sec = 120 ,
170+ timeout_sec = max ( 15 , int ( timeout_sec )) ,
143171 )
144172 if result .return_code != 0 :
145173 return {
@@ -166,7 +194,11 @@ def _usage_value(value: Any, name: str) -> int:
166194 return int (result or 0 )
167195
168196 def _completion_kwargs (
169- self , messages : list [dict [str , Any ]], * , tools : bool = True
197+ self ,
198+ messages : list [dict [str , Any ]],
199+ * ,
200+ tools : bool = True ,
201+ max_tokens : int = 12_000 ,
170202 ) -> dict [str , Any ]:
171203 # Reasoning models replaced max_tokens with max_completion_tokens and
172204 # reject the old name outright ("Unsupported parameter: 'max_tokens'
@@ -181,14 +213,43 @@ def _completion_kwargs(
181213 "model" : self ._api_model ,
182214 "messages" : messages ,
183215 }
184- kwargs [_token_limit_key ] = 12_000
216+ kwargs [_token_limit_key ] = max_tokens
185217 if tools :
186218 kwargs ["tools" ] = TOOLS
187219 # Capability, not provider: see _is_reasoning_model.
188220 if _is_reasoning_model (self ._api_model ):
189221 kwargs ["reasoning_effort" ] = "medium"
190222 return kwargs
191223
224+ async def _complete (
225+ self ,
226+ messages : list [dict [str , Any ]],
227+ * ,
228+ tools : bool = True ,
229+ timeout : float ,
230+ max_retries : int | None = None ,
231+ max_tokens : int = 12_000 ,
232+ ) -> Any :
233+ """One chat completion, validated for shape.
234+
235+ A gateway can answer 200 with a body the SDK does not turn into a
236+ completion (observed once as `'str' object has no attribute 'usage'`),
237+ so anything without usable `.choices` is rejected here rather than
238+ crashing the trial on an attribute access two lines later.
239+ """
240+ options : dict [str , Any ] = {"timeout" : max (30.0 , timeout )}
241+ if max_retries is not None :
242+ options ["max_retries" ] = max_retries
243+ response = await self ._client .with_options (** options ).chat .completions .create (
244+ ** self ._completion_kwargs (messages , tools = tools , max_tokens = max_tokens )
245+ )
246+ choices = getattr (response , "choices" , None )
247+ if not choices or getattr (choices [0 ], "message" , None ) is None :
248+ raise RuntimeError (
249+ f"malformed completion from provider: { type (response ).__name__ } "
250+ )
251+ return response
252+
192253 def _account (self , usage : Any , totals : dict [str , int ]) -> None :
193254 totals ["input" ] += self ._usage_value (usage , "prompt_tokens" )
194255 totals ["output" ] += self ._usage_value (usage , "completion_tokens" )
@@ -211,25 +272,50 @@ async def run(
211272 {"role" : "user" , "content" : instruction },
212273 ]
213274 totals = {"input" : 0 , "output" : 0 , "cached" : 0 }
275+ deadline = time .monotonic () + AGENT_BUDGET_SEC
276+
277+ def remaining () -> float :
278+ return deadline - time .monotonic ()
279+
280+ completed = False
281+ out_of_time = False
282+ out_of_answers = False
283+ empty_turns = 0
284+ last_turn = 0
214285
215286 for turn in range (1 , MAX_TURNS + 1 ):
216- response = await self ._client .chat .completions .create (
217- ** self ._completion_kwargs (messages )
287+ last_turn = turn
288+ # Stop starting new work once only the final-answer reserve is left,
289+ # so the answer is written before harbor's AgentTimeoutError fires.
290+ if remaining () <= FINAL_RESERVE_SEC :
291+ out_of_time = True
292+ self ._trace ({"turn" : turn , "stopping" : "deadline" , "left" : remaining ()})
293+ break
294+ call_started = time .monotonic ()
295+ response = await self ._complete (
296+ messages , timeout = remaining () - FINAL_RESERVE_SEC / 2
218297 )
219- self ._account (response .usage , totals )
298+ llm_sec = time .monotonic () - call_started
299+ self ._account (getattr (response , "usage" , None ), totals )
220300 message = response .choices [0 ].message
301+ finish_reason = getattr (response .choices [0 ], "finish_reason" , None )
221302 calls = message .tool_calls or []
222303 self ._trace (
223304 {
224305 "turn" : turn ,
225306 "content" : message .content ,
307+ "finish_reason" : finish_reason ,
308+ "llm_sec" : round (llm_sec , 2 ),
309+ "left_sec" : round (remaining (), 1 ),
226310 "tool_calls" : [
227311 {"name" : c .function .name , "arguments" : c .function .arguments }
228312 for c in calls
229313 ],
230314 }
231315 )
232316 # Record the assistant turn (with any tool calls) in the history.
317+ # An assistant message carrying neither is rejected by strict
318+ # OpenAI-compatible servers, so it is dropped rather than appended.
233319 assistant : dict [str , Any ] = {"role" : "assistant" }
234320 if message .content :
235321 assistant ["content" ] = message .content
@@ -245,44 +331,89 @@ async def run(
245331 }
246332 for c in calls
247333 ]
248- messages .append (assistant )
334+ if len (assistant ) > 1 :
335+ messages .append (assistant )
249336
250337 if not calls :
251- if (message .content or "" ).strip ():
338+ # `length` means the model was cut off mid-thought, so its
339+ # content is truncated reasoning, not an answer. This target
340+ # writes its chain-of-thought into ordinary `content` (it has no
341+ # separate reasoning_content channel) and can run past the token
342+ # cap, and the truncated ramble was previously submitted verbatim
343+ # as the final answer, which the judge always scores 0. Ask for
344+ # the answer instead of submitting the thinking.
345+ truncated = finish_reason == "length"
346+ if (message .content or "" ).strip () and not truncated :
252347 await self ._submit (environment , message .content )
348+ completed = True
253349 context .metadata = {
254350 "turns" : turn ,
255351 "trace" : "browsecomp-plus-trace.jsonl" ,
256352 }
257353 break
258- raise RuntimeError ("model returned neither a response nor a tool call" )
354+ empty_turns += 1
355+ self ._trace (
356+ {"turn" : turn , "recover" : empty_turns , "truncated" : truncated }
357+ )
358+ if empty_turns >= MAX_EMPTY_RETRIES :
359+ out_of_answers = True
360+ break
361+ messages .append (
362+ {
363+ "role" : "user" ,
364+ "content" : (
365+ "You were cut off mid-thought. Stop reasoning and "
366+ "reply with only the three required labeled lines "
367+ "(Explanation, Exact Answer, Confidence), keeping the "
368+ "explanation to a few sentences."
369+ if truncated
370+ else "Your last message was empty. Either call a tool "
371+ "or give your final response in the required format."
372+ ),
373+ }
374+ )
375+ continue
376+ empty_turns = 0
259377
260378 submitted = False
261379 for call in calls :
262- try :
263- arguments = json .loads (call .function .arguments )
264- if call .function .name == "search" :
265- result = await self ._index_command (
266- environment , "search" , arguments ["query" ]
267- )
268- elif call .function .name == "get_document" :
269- result = await self ._index_command (
270- environment , "get-document" , arguments ["docid" ]
271- )
272- elif call .function .name == "submit_response" :
273- await self ._submit (environment , arguments ["response" ])
274- result = {"submitted" : True }
275- submitted = True
276- else :
277- result = {"error" : f"unknown tool: { call .function .name } " }
278- except (
279- json .JSONDecodeError ,
280- KeyError ,
281- OSError ,
282- TypeError ,
283- ValueError ,
284- ) as error :
285- result = {"error" : str (error )}
380+ # Every tool_call must still get a tool message even once the
381+ # deadline hits, or the forced-final request below is rejected
382+ # for an unanswered tool call. So stub, never skip.
383+ if out_of_time or remaining () <= FINAL_RESERVE_SEC :
384+ if not out_of_time :
385+ out_of_time = True
386+ self ._trace ({"turn" : turn , "stopping" : "deadline-midturn" })
387+ result = {"error" : "search budget exhausted" }
388+ else :
389+ try :
390+ arguments = json .loads (call .function .arguments )
391+ budget = remaining () - FINAL_RESERVE_SEC
392+ if call .function .name == "search" :
393+ result = await self ._index_command (
394+ environment , "search" , arguments ["query" ], budget
395+ )
396+ elif call .function .name == "get_document" :
397+ result = await self ._index_command (
398+ environment , "get-document" , arguments ["docid" ], budget
399+ )
400+ elif call .function .name == "submit_response" :
401+ await self ._submit (environment , arguments ["response" ])
402+ result = {"submitted" : True }
403+ submitted = True
404+ else :
405+ result = {"error" : f"unknown tool: { call .function .name } " }
406+ except (
407+ json .JSONDecodeError ,
408+ KeyError ,
409+ OSError ,
410+ TypeError ,
411+ ValueError ,
412+ ) as error :
413+ result = {"error" : str (error )}
414+ # Traced on both paths. A stub is precisely what a post-hoc audit
415+ # needs to see (which searches were abandoned to the clock), and
416+ # tracing only the executed branch left it invisible.
286417 self ._trace (
287418 {"turn" : turn , "tool" : call .function .name , "result" : result }
288419 )
@@ -298,39 +429,81 @@ async def run(
298429 if submitted :
299430 break
300431 if submitted :
432+ completed = True
301433 context .metadata = {
302434 "turns" : turn ,
303435 "trace" : "browsecomp-plus-trace.jsonl" ,
304436 }
305437 break
306- else :
307- # Turn budget exhausted: force one final tool-free response from what
308- # was retrieved rather than crashing, so the case scores best-effort
438+ if out_of_time :
439+ break
440+
441+ if not completed :
442+ # Budget exhausted (wall clock, turns, or a model that kept coming
443+ # back empty): force one final tool-free response from what was
444+ # retrieved rather than crashing, so the case scores best-effort
309445 # instead of being lost with no answer recorded.
446+ stub = (
447+ "Explanation: no answer could be determined within the search "
448+ "budget.\n Exact Answer: unknown\n Confidence: 0%"
449+ )
450+ # Three distinct exits land here and a post-hoc scorer has to tell
451+ # them apart: the wall clock, repeated truncated/empty turns, and
452+ # plain MAX_TURNS exhaustion.
453+ reason = (
454+ "deadline"
455+ if out_of_time
456+ else "empty_retries"
457+ if out_of_answers
458+ else "turns"
459+ )
310460 messages .append (
311461 {
312462 "role" : "user" ,
313463 "content" : (
314- "You have used your full search budget. Give your single "
315- "best final response now, in the required format, based on "
316- "the documents you have retrieved."
464+ "You have used your full search budget. Do not reason "
465+ "further. Reply with only the three required labeled "
466+ "lines (Explanation, Exact Answer, Confidence), keeping "
467+ "the explanation to a few sentences, giving your single "
468+ "best answer from the documents you have retrieved."
317469 ),
318470 }
319471 )
320- final = await self ._client .chat .completions .create (
321- ** self ._completion_kwargs (messages , tools = False )
322- )
323- self ._account (final .usage , totals )
324- answer = (final .choices [0 ].message .content or "" ).strip () or (
325- "Explanation: no answer could be determined within the search "
326- "budget.\n Exact Answer: unknown\n Confidence: 0%"
327- )
472+ # This call is the last thing standing between a completed trial and
473+ # a lost one, so a failure here degrades to the stub answer instead
474+ # of propagating.
475+ answer = stub
476+ try :
477+ # Bounded retries: the SDK default of 8 could outlast the
478+ # reserve on its own and hand the trial back to the killer.
479+ # A tight token cap keeps the model from spending the whole
480+ # reserve on chain-of-thought and getting truncated again.
481+ final = await self ._complete (
482+ messages ,
483+ tools = False ,
484+ timeout = max (30.0 , remaining () - 20 ),
485+ max_retries = 1 ,
486+ max_tokens = ANSWER_MAX_TOKENS ,
487+ )
488+ self ._account (getattr (final , "usage" , None ), totals )
489+ answer = (final .choices [0 ].message .content or "" ).strip () or stub
490+ except Exception as error : # noqa: BLE001 - last-resort safety net
491+ self ._trace ({"forced_final_error" : f"{ type (error ).__name__ } : { error } " })
328492 await self ._submit (environment , answer )
329- self ._trace ({"turn" : MAX_TURNS , "forced_final_answer" : answer })
493+ self ._trace (
494+ {
495+ "turn" : last_turn ,
496+ "forced_final_answer" : answer ,
497+ "reason" : reason ,
498+ }
499+ )
330500 context .metadata = {
331- "turns" : MAX_TURNS ,
501+ "turns" : last_turn ,
332502 "trace" : "browsecomp-plus-trace.jsonl" ,
333503 "forced_final" : True ,
504+ # Lets a post-hoc conservative score treat exactly the trials
505+ # that the pre-fix agent would have lost as zeros.
506+ "forced_final_reason" : reason ,
334507 }
335508
336509 context .n_input_tokens = totals ["input" ]
0 commit comments