55import json
66import os
77import shlex
8+ import time
89from typing import Any , override
910
1011from harbor .agents .base import BaseAgent
@@ -25,6 +26,28 @@ def _is_reasoning_model(model: str) -> bool:
2526MAX_TURNS = 32
2627MAX_TOOL_OUTPUT_CHARS = 40_000
2728
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+
2851INSTRUCTIONS = """You are a persistent deep-research agent working only against the
2952fixed BrowseComp-Plus corpus. Find the precise answer by issuing focused searches,
3053opening promising documents, reformulating queries, and cross-checking evidence.
@@ -147,13 +170,17 @@ def _trace(self, event: dict[str, Any]) -> None:
147170 file .write (json .dumps (event , ensure_ascii = False , default = str ) + "\n " )
148171
149172 async def _index_command (
150- self , environment : BaseEnvironment , command : str , value : str
173+ self ,
174+ environment : BaseEnvironment ,
175+ command : str ,
176+ value : str ,
177+ timeout_sec : float = 120 ,
151178 ) -> dict [str , Any ]:
152179 flag = "--query" if command == "search" else "--docid"
153180 result = await environment .exec (
154181 f"python3 /opt/browsecomp/search.py { command } { flag } { shlex .quote (value )} " ,
155182 cwd = "/app" ,
156- timeout_sec = 120 ,
183+ timeout_sec = max ( 15 , int ( timeout_sec )) ,
157184 )
158185 if result .return_code != 0 :
159186 return {
@@ -180,7 +207,11 @@ def _usage_value(value: Any, name: str) -> int:
180207 return int (result or 0 )
181208
182209 def _completion_kwargs (
183- self , messages : list [dict [str , Any ]], * , tools : bool = True
210+ self ,
211+ messages : list [dict [str , Any ]],
212+ * ,
213+ tools : bool = True ,
214+ max_tokens : int = 12_000 ,
184215 ) -> dict [str , Any ]:
185216 # Reasoning models replaced max_tokens with max_completion_tokens and
186217 # reject the old name outright ("Unsupported parameter: 'max_tokens'
@@ -195,14 +226,43 @@ def _completion_kwargs(
195226 "model" : self ._api_model ,
196227 "messages" : messages ,
197228 }
198- kwargs [_token_limit_key ] = 12_000
229+ kwargs [_token_limit_key ] = max_tokens
199230 if tools :
200231 kwargs ["tools" ] = TOOLS
201232 # Capability, not provider: see _is_reasoning_model.
202233 if _is_reasoning_model (self ._api_model ):
203234 kwargs ["reasoning_effort" ] = "medium"
204235 return kwargs
205236
237+ async def _complete (
238+ self ,
239+ messages : list [dict [str , Any ]],
240+ * ,
241+ tools : bool = True ,
242+ timeout : float ,
243+ max_retries : int | None = None ,
244+ max_tokens : int = 12_000 ,
245+ ) -> Any :
246+ """One chat completion, validated for shape.
247+
248+ A gateway can answer 200 with a body the SDK does not turn into a
249+ completion (observed once as `'str' object has no attribute 'usage'`),
250+ so anything without usable `.choices` is rejected here rather than
251+ crashing the trial on an attribute access two lines later.
252+ """
253+ options : dict [str , Any ] = {"timeout" : max (30.0 , timeout )}
254+ if max_retries is not None :
255+ options ["max_retries" ] = max_retries
256+ response = await self ._client .with_options (** options ).chat .completions .create (
257+ ** self ._completion_kwargs (messages , tools = tools , max_tokens = max_tokens )
258+ )
259+ choices = getattr (response , "choices" , None )
260+ if not choices or getattr (choices [0 ], "message" , None ) is None :
261+ raise RuntimeError (
262+ f"malformed completion from provider: { type (response ).__name__ } "
263+ )
264+ return response
265+
206266 def _account (self , usage : Any , totals : dict [str , int ]) -> None :
207267 totals ["input" ] += self ._usage_value (usage , "prompt_tokens" )
208268 totals ["output" ] += self ._usage_value (usage , "completion_tokens" )
@@ -225,25 +285,50 @@ async def run(
225285 {"role" : "user" , "content" : instruction },
226286 ]
227287 totals = {"input" : 0 , "output" : 0 , "cached" : 0 }
288+ deadline = time .monotonic () + AGENT_BUDGET_SEC
289+
290+ def remaining () -> float :
291+ return deadline - time .monotonic ()
292+
293+ completed = False
294+ out_of_time = False
295+ out_of_answers = False
296+ empty_turns = 0
297+ last_turn = 0
228298
229299 for turn in range (1 , MAX_TURNS + 1 ):
230- response = await self ._client .chat .completions .create (
231- ** self ._completion_kwargs (messages )
300+ last_turn = turn
301+ # Stop starting new work once only the final-answer reserve is left,
302+ # so the answer is written before harbor's AgentTimeoutError fires.
303+ if remaining () <= FINAL_RESERVE_SEC :
304+ out_of_time = True
305+ self ._trace ({"turn" : turn , "stopping" : "deadline" , "left" : remaining ()})
306+ break
307+ call_started = time .monotonic ()
308+ response = await self ._complete (
309+ messages , timeout = remaining () - FINAL_RESERVE_SEC / 2
232310 )
233- self ._account (response .usage , totals )
311+ llm_sec = time .monotonic () - call_started
312+ self ._account (getattr (response , "usage" , None ), totals )
234313 message = response .choices [0 ].message
314+ finish_reason = getattr (response .choices [0 ], "finish_reason" , None )
235315 calls = message .tool_calls or []
236316 self ._trace (
237317 {
238318 "turn" : turn ,
239319 "content" : message .content ,
320+ "finish_reason" : finish_reason ,
321+ "llm_sec" : round (llm_sec , 2 ),
322+ "left_sec" : round (remaining (), 1 ),
240323 "tool_calls" : [
241324 {"name" : c .function .name , "arguments" : c .function .arguments }
242325 for c in calls
243326 ],
244327 }
245328 )
246329 # Record the assistant turn (with any tool calls) in the history.
330+ # An assistant message carrying neither is rejected by strict
331+ # OpenAI-compatible servers, so it is dropped rather than appended.
247332 assistant : dict [str , Any ] = {"role" : "assistant" }
248333 if message .content :
249334 assistant ["content" ] = message .content
@@ -259,44 +344,89 @@ async def run(
259344 }
260345 for c in calls
261346 ]
262- messages .append (assistant )
347+ if len (assistant ) > 1 :
348+ messages .append (assistant )
263349
264350 if not calls :
265- if (message .content or "" ).strip ():
351+ # `length` means the model was cut off mid-thought, so its
352+ # content is truncated reasoning, not an answer. This target
353+ # writes its chain-of-thought into ordinary `content` (it has no
354+ # separate reasoning_content channel) and can run past the token
355+ # cap, and the truncated ramble was previously submitted verbatim
356+ # as the final answer, which the judge always scores 0. Ask for
357+ # the answer instead of submitting the thinking.
358+ truncated = finish_reason == "length"
359+ if (message .content or "" ).strip () and not truncated :
266360 await self ._submit (environment , message .content )
361+ completed = True
267362 context .metadata = {
268363 "turns" : turn ,
269364 "trace" : "browsecomp-plus-trace.jsonl" ,
270365 }
271366 break
272- raise RuntimeError ("model returned neither a response nor a tool call" )
367+ empty_turns += 1
368+ self ._trace (
369+ {"turn" : turn , "recover" : empty_turns , "truncated" : truncated }
370+ )
371+ if empty_turns >= MAX_EMPTY_RETRIES :
372+ out_of_answers = True
373+ break
374+ messages .append (
375+ {
376+ "role" : "user" ,
377+ "content" : (
378+ "You were cut off mid-thought. Stop reasoning and "
379+ "reply with only the three required labeled lines "
380+ "(Explanation, Exact Answer, Confidence), keeping the "
381+ "explanation to a few sentences."
382+ if truncated
383+ else "Your last message was empty. Either call a tool "
384+ "or give your final response in the required format."
385+ ),
386+ }
387+ )
388+ continue
389+ empty_turns = 0
273390
274391 submitted = False
275392 for call in calls :
276- try :
277- arguments = json .loads (call .function .arguments )
278- if call .function .name == "search" :
279- result = await self ._index_command (
280- environment , "search" , arguments ["query" ]
281- )
282- elif call .function .name == "get_document" :
283- result = await self ._index_command (
284- environment , "get-document" , arguments ["docid" ]
285- )
286- elif call .function .name == "submit_response" :
287- await self ._submit (environment , arguments ["response" ])
288- result = {"submitted" : True }
289- submitted = True
290- else :
291- result = {"error" : f"unknown tool: { call .function .name } " }
292- except (
293- json .JSONDecodeError ,
294- KeyError ,
295- OSError ,
296- TypeError ,
297- ValueError ,
298- ) as error :
299- result = {"error" : str (error )}
393+ # Every tool_call must still get a tool message even once the
394+ # deadline hits, or the forced-final request below is rejected
395+ # for an unanswered tool call. So stub, never skip.
396+ if out_of_time or remaining () <= FINAL_RESERVE_SEC :
397+ if not out_of_time :
398+ out_of_time = True
399+ self ._trace ({"turn" : turn , "stopping" : "deadline-midturn" })
400+ result = {"error" : "search budget exhausted" }
401+ else :
402+ try :
403+ arguments = json .loads (call .function .arguments )
404+ budget = remaining () - FINAL_RESERVE_SEC
405+ if call .function .name == "search" :
406+ result = await self ._index_command (
407+ environment , "search" , arguments ["query" ], budget
408+ )
409+ elif call .function .name == "get_document" :
410+ result = await self ._index_command (
411+ environment , "get-document" , arguments ["docid" ], budget
412+ )
413+ elif call .function .name == "submit_response" :
414+ await self ._submit (environment , arguments ["response" ])
415+ result = {"submitted" : True }
416+ submitted = True
417+ else :
418+ result = {"error" : f"unknown tool: { call .function .name } " }
419+ except (
420+ json .JSONDecodeError ,
421+ KeyError ,
422+ OSError ,
423+ TypeError ,
424+ ValueError ,
425+ ) as error :
426+ result = {"error" : str (error )}
427+ # Traced on both paths. A stub is precisely what a post-hoc audit
428+ # needs to see (which searches were abandoned to the clock), and
429+ # tracing only the executed branch left it invisible.
300430 self ._trace (
301431 {"turn" : turn , "tool" : call .function .name , "result" : result }
302432 )
@@ -312,39 +442,81 @@ async def run(
312442 if submitted :
313443 break
314444 if submitted :
445+ completed = True
315446 context .metadata = {
316447 "turns" : turn ,
317448 "trace" : "browsecomp-plus-trace.jsonl" ,
318449 }
319450 break
320- else :
321- # Turn budget exhausted: force one final tool-free response from what
322- # was retrieved rather than crashing, so the case scores best-effort
451+ if out_of_time :
452+ break
453+
454+ if not completed :
455+ # Budget exhausted (wall clock, turns, or a model that kept coming
456+ # back empty): force one final tool-free response from what was
457+ # retrieved rather than crashing, so the case scores best-effort
323458 # instead of being lost with no answer recorded.
459+ stub = (
460+ "Explanation: no answer could be determined within the search "
461+ "budget.\n Exact Answer: unknown\n Confidence: 0%"
462+ )
463+ # Three distinct exits land here and a post-hoc scorer has to tell
464+ # them apart: the wall clock, repeated truncated/empty turns, and
465+ # plain MAX_TURNS exhaustion.
466+ reason = (
467+ "deadline"
468+ if out_of_time
469+ else "empty_retries"
470+ if out_of_answers
471+ else "turns"
472+ )
324473 messages .append (
325474 {
326475 "role" : "user" ,
327476 "content" : (
328- "You have used your full search budget. Give your single "
329- "best final response now, in the required format, based on "
330- "the documents you have retrieved."
477+ "You have used your full search budget. Do not reason "
478+ "further. Reply with only the three required labeled "
479+ "lines (Explanation, Exact Answer, Confidence), keeping "
480+ "the explanation to a few sentences, giving your single "
481+ "best answer from the documents you have retrieved."
331482 ),
332483 }
333484 )
334- final = await self ._client .chat .completions .create (
335- ** self ._completion_kwargs (messages , tools = False )
336- )
337- self ._account (final .usage , totals )
338- answer = (final .choices [0 ].message .content or "" ).strip () or (
339- "Explanation: no answer could be determined within the search "
340- "budget.\n Exact Answer: unknown\n Confidence: 0%"
341- )
485+ # This call is the last thing standing between a completed trial and
486+ # a lost one, so a failure here degrades to the stub answer instead
487+ # of propagating.
488+ answer = stub
489+ try :
490+ # Bounded retries: the SDK default of 8 could outlast the
491+ # reserve on its own and hand the trial back to the killer.
492+ # A tight token cap keeps the model from spending the whole
493+ # reserve on chain-of-thought and getting truncated again.
494+ final = await self ._complete (
495+ messages ,
496+ tools = False ,
497+ timeout = max (30.0 , remaining () - 20 ),
498+ max_retries = 1 ,
499+ max_tokens = ANSWER_MAX_TOKENS ,
500+ )
501+ self ._account (getattr (final , "usage" , None ), totals )
502+ answer = (final .choices [0 ].message .content or "" ).strip () or stub
503+ except Exception as error : # noqa: BLE001 - last-resort safety net
504+ self ._trace ({"forced_final_error" : f"{ type (error ).__name__ } : { error } " })
342505 await self ._submit (environment , answer )
343- self ._trace ({"turn" : MAX_TURNS , "forced_final_answer" : answer })
506+ self ._trace (
507+ {
508+ "turn" : last_turn ,
509+ "forced_final_answer" : answer ,
510+ "reason" : reason ,
511+ }
512+ )
344513 context .metadata = {
345- "turns" : MAX_TURNS ,
514+ "turns" : last_turn ,
346515 "trace" : "browsecomp-plus-trace.jsonl" ,
347516 "forced_final" : True ,
517+ # Lets a post-hoc conservative score treat exactly the trials
518+ # that the pre-fix agent would have lost as zeros.
519+ "forced_final_reason" : reason ,
348520 }
349521
350522 context .n_input_tokens = totals ["input" ]
0 commit comments