import contextlib

import itertools

import types

import typing

from collections.abc import Callable, Sequence

from typing import Any, Literal, TypedDict

import inspect_ai.agent

import inspect_ai.log

import inspect_ai.model

import inspect_ai.model._generate_config

import inspect_ai.solver

import inspect_ai.tool

import inspect_ai.util

import pydantic

import shortuuid

from inspect_ai._util.notgiven import NotGiven

import metr_agents.tools

ANSWER_DELIMITER = "sep_TFLTJ88PEK"

NOT_GIVEN = NotGiven()

# Checkpointer key for default_generator's persisted original input.

INITIAL_MESSAGES_CHECKPOINT_KEY = "metr_agents_initial_messages"

# Checkpointer key inspect's `compaction()` registers its state under. Must

# match the literal in inspect_ai.model._compaction._compaction.

COMPACTION_CHECKPOINT_KEY = "compaction"

# Injected once when a sample resumes after a crash (checkpointer attempt

# "resume"). When inspect-ai exposes Task.on_resume / ResumeReport (PR #4383),

# this can additionally be suppressed for transparent resumes and extended with

# task-supplied detail.

RESUME_NOTICE = (

"System note: resumed from a checkpoint after a crash. Your memory of the "

"run was kept, but the environment may not match it. It could be unchanged, "

"partly reset, or wiped. The checkpoint was taken one or more steps before "

"the crash happened. Crashes are often caused by running out of memory or "

"disk space, so be careful with operations that use a lot of memory or write "

"large files to disk."

)

def format_duration(seconds: float) -> str:

total_seconds = round(seconds)

hours, remainder = divmod(total_seconds, 3600)

minutes, secs = divmod(remainder, 60)

parts: list[str] = []

if hours:

parts.append(f"{hours} {'hour' if hours == 1 else 'hours'}")

if minutes:

parts.append(f"{minutes} {'minute' if minutes == 1 else 'minutes'}")

if secs or not parts:

parts.append(f"{secs} {'second' if secs == 1 else 'seconds'}")

if len(parts) == 1:

return parts[0]

return ", ".join(parts[:-1]) + " and " + parts[-1]

class LimitWarningConfig(pydantic.BaseModel):

threshold: float = pydantic.Field(ge=0.0, le=1.0)

message: str

class LimitWarningDefaults(pydantic.BaseModel):

near: LimitWarningConfig | None = LimitWarningConfig(

threshold=0.80,

message="You are nearing the limit. Prepare to submit your answer.",

)

very_near: LimitWarningConfig | None = LimitWarningConfig(

threshold=0.95,

message="You are almost at the limit. Submit your answer as soon as possible.",

)

@pydantic.model_validator(mode="after")

def _check_threshold_ordering(self) -> "LimitWarningDefaults":

if (

self.near is not None

and self.very_near is not None

and self.near.threshold >= self.very_near.threshold

):

raise ValueError(

f"near threshold ({self.near.threshold}) must be less than "

+ f"very_near threshold ({self.very_near.threshold})"

)

return self

class LimitTypeConfig(pydantic.BaseModel):

message_template: str

near: LimitWarningConfig | None | Literal["default"] = "default"

very_near: LimitWarningConfig | None | Literal["default"] = "default"

class LimitMessageConfig(pydantic.BaseModel):

defaults: LimitWarningDefaults = LimitWarningDefaults()

token: LimitTypeConfig | None = LimitTypeConfig(

message_template=(

"Tokens used so far: {usage}/{limit}. "

+ "You have used {pct_used:.2f}% of your total token budget."

),

)

time: LimitTypeConfig | None = LimitTypeConfig(

message_template=(

"Time elapsed: {usage} of {limit} available. "

+ "You have used {pct_used:.2f}% of your total time budget."

),

)

working: LimitTypeConfig | None = LimitTypeConfig(

message_template=(

"Working time elapsed: {usage} of {limit} available. "

+ "You have used {pct_used:.2f}% of your total working time budget."

),

)

cost: LimitTypeConfig | None = LimitTypeConfig(

message_template=(

"Token cost so far: ${usage} of ${limit} available. "

+ "You have used {pct_used:.2f}% of your total token cost budget."

),

)

def _resolve_warning(

type_config: LimitTypeConfig,

defaults: LimitWarningDefaults,

field: Literal["near", "very_near"],

) -> LimitWarningConfig | None:

value = getattr(type_config, field)

if value == "default":

return getattr(defaults, field)

return value

def format_limit_usage_message(

usage: float,

limit: float,

type_config: LimitTypeConfig,

defaults: LimitWarningDefaults,

format_value: Callable[[float], str],

) -> str:

relative_usage = usage / limit if limit > 0 else 0

try:

message = type_config.message_template.format(

usage=format_value(usage),

limit=format_value(limit),

usage_ratio=relative_usage,

pct_used=relative_usage * 100,

)

except KeyError as e:

raise ValueError(

f"Unknown placeholder {e} in message_template. "

+ "Available placeholders: {usage}, {limit}, {usage_ratio}, {pct_used}"

) from e

if limit > 0:

very_near = _resolve_warning(type_config, defaults, "very_near")

near = _resolve_warning(type_config, defaults, "near")

if very_near is not None and relative_usage > very_near.threshold:

message += f" {very_near.message}"

elif near is not None and relative_usage > near.threshold:

message += f" {near.message}"

return message

LIMIT_FORMAT_VALUES: types.MappingProxyType[str, Callable[[float], str]] = (

types.MappingProxyType(

{

"token": lambda v: f"{int(v):,}",

"time": format_duration,

"working": format_duration,

"cost": lambda v: f"{v:.4f}",

}

)

)

def collect_limit_messages(

config: LimitMessageConfig,

sample_limits: inspect_ai.util.SampleLimits,

) -> list[str]:

messages: list[str] = []

for limit_type, limit_formatter in LIMIT_FORMAT_VALUES.items():

type_config: LimitTypeConfig | None = getattr(config, limit_type)

if type_config is None:

continue

limit_info = getattr(sample_limits, limit_type)

if limit_info.limit is None:

continue

if float(limit_info.usage) == 0:

continue

messages.append(

format_limit_usage_message(

usage=float(limit_info.usage),

limit=float(limit_info.limit),

type_config=type_config,

defaults=config.defaults,

format_value=limit_formatter,

)

)

if messages:

messages[0] = f"[usage] {messages[0]}"

return messages

DEFAULT_PROCEED_PROMPT = (

"\nPlease proceed to the next step using your best judgement. If you believe you\n"

"have completed the task, please call the `submit()` tool.\n"

)

@inspect_ai.agent.agent

def limit_usage_message(

proceed_prompt: str | None = DEFAULT_PROCEED_PROMPT,

config: LimitMessageConfig | None = None,

) -> inspect_ai.agent.Agent:

if config is None:

config = LimitMessageConfig()

async def execute(

state: inspect_ai.agent.AgentState,

) -> inspect_ai.agent.AgentState:

main_message = (

""

if state.output.message.tool_calls or proceed_prompt is None

else proceed_prompt

)

if main_message:

state.messages.append(

inspect_ai.model.ChatMessageUser(content=main_message)

)

for msg in collect_limit_messages(config, inspect_ai.util.sample_limits()):

state.messages.append(inspect_ai.model.ChatMessageUser(content=msg))

return state

return execute

DEFAULT_EARLY_SUBMIT_MESSAGE = (

"You are trying to submit but have only used {pct_used:.0f}% of your token budget "

"({token_usage:,}/{token_limit:,}). You can only submit after you have used at least "

"{pct_required:.0f}% of your token budget. Continue working until you have used at "

"least {pct_required:.0f}% of your token budget."

)

COMPACTION_CLASSES = types.MappingProxyType(

{

"auto": inspect_ai.model.CompactionAuto,

"edit": inspect_ai.model.CompactionEdit,

"native": inspect_ai.model.CompactionNative,

"summary": inspect_ai.model.CompactionSummary,

"trim": inspect_ai.model.CompactionTrim,

}

)

class OptionalReactKwargs(TypedDict, total=False):

prompt: str | inspect_ai.agent.AgentPrompt | None

class CompactionConfig(pydantic.BaseModel):

strategy: str

args: dict[str, Any] = pydantic.Field(default_factory=dict)

class RebindableCheckpointer:

"""Checkpointer facade that lets a tracked key be re-registered.

Inspect's checkpointer raises if a key is tracked twice, but we rebuild the

compaction handler on every handoff (so compaction doesn't judge the new

context by the previous agent's history) and each rebuild calls

`track("compaction", ...)`.

We register once with the real checkpointer, handing it a shim callback that

forwards to whichever callback we saw last, and return `initial_value` for

later registrations. So a checkpoint fire captures the live handler's state

rather than a stale one's, and each rebuild starts from fresh state.

`discard_checkpointed` additionally drops a *restored* value: after a

handoff, any checkpointed compaction state describes the context we just

cleared, and replaying it would resurrect the previous agent's messages.

"""

_checkpointer: inspect_ai.util.Checkpointer

_callbacks: dict[str, Callable[[], object]]

_value_types: dict[str, object]

_discarded: set[str]

def __init__(self, checkpointer: inspect_ai.util.Checkpointer) -> None:

self._checkpointer = checkpointer

self._callbacks = {}

self._value_types = {}

self._discarded = set()

def discard_checkpointed(self, key: str) -> None:

"""Ignore any checkpointed value restored for `key` when it registers."""

self._discarded.add(key)

def track[T](

self,

key: str,

callback: Callable[[], T],

initial_value: T,

*,

value_type: type[T] | None = None,

) -> T:

if key in self._callbacks:

if value_type is not self._value_types[key]:

raise ValueError(

f"track({key!r}) re-registered with a different value_type; "

+ "the underlying checkpointer keeps the first registration's "

+ "type, so resume would deserialize into the wrong shape."

)

self._callbacks[key] = callback

return initial_value

self._callbacks[key] = callback

self._value_types[key] = value_type

# The shim is `Callable[[], object]` because `_callbacks` is

# heterogeneous; each key is only ever re-registered with the same `T`,

# which the `value_type` check above enforces.

shim = typing.cast("Callable[[], T]", lambda: self._callbacks[key]())

restored = self._checkpointer.track(

key, shim, initial_value, value_type=value_type

)

return initial_value if key in self._discarded else restored

@property

def attempt(self) -> Literal["initial", "resume", "resume_for_scoring"]:

return self._checkpointer.attempt

async def tick(self) -> None:

await self._checkpointer.tick()

async def checkpoint(self) -> None:

await self._checkpointer.checkpoint()

def span_session(self) -> contextlib.AbstractAsyncContextManager[None]:

return self._checkpointer.span_session()

def build_compact_handler(

compaction: CompactionConfig | dict[str, str | dict[str, Any]],

initial_messages: list[inspect_ai.model.ChatMessage] | None,

tools: Sequence[inspect_ai.tool.Tool] | None,

checkpointer: inspect_ai.util.Checkpointer | None = None,

) -> inspect_ai.model.Compact:

compact_config = CompactionConfig.model_validate(compaction)

if compact_config.strategy not in COMPACTION_CLASSES:

raise ValueError(

f"Unknown compaction strategy: {compact_config.strategy!r}. "

+ f"Valid options are: {', '.join(COMPACTION_CLASSES.keys())}"

)

compact_class = COMPACTION_CLASSES[compact_config.strategy]

strategy = compact_class(**compact_config.args)

prefix = initial_messages or []

# Pass the checkpointer only when present so compaction state is

# tracked/restored; omit it otherwise (inspect's no-op default applies).

if checkpointer is None:

return inspect_ai.model.compaction(strategy, prefix=prefix, tools=tools)

return inspect_ai.model.compaction(

strategy, prefix=prefix, tools=tools, checkpointer=checkpointer

)

@inspect_ai.agent.agent

def default_generator(

compaction: CompactionConfig | None = None,

retry_refusals: int | None = None,

handoff: bool = False,

) -> inspect_ai.agent.Agent:

"""Default generator function for the react agent.

Args:

compaction: Compaction config, or None to disable compaction.

retry_refusals: How many times to retry a `content_filter` refusal.

handoff: Whether to honour handoffs left by `metr_agents.tools.handoff_submit`

or `metr_agents.tools.handoff`.

Only `react_with_handoff_submit` sets this; other agents must not read the

handoff store keys.

"""

_compact: inspect_ai.model.Compact | None = None

_initial_messages: list[inspect_ai.model.ChatMessage] | None = None

_input_reminder_content: list[inspect_ai.model.Content] | None = None

_tools: Sequence[inspect_ai.tool.Tool] | None = None

_resume_notice_done = False

_rebindable_checkpointer: RebindableCheckpointer | None = None

async def execute(

state: inspect_ai.agent.AgentState,

tools: Sequence[inspect_ai.tool.Tool],

) -> inspect_ai.agent.AgentState:

nonlocal _compact, _initial_messages, _input_reminder_content, _tools

nonlocal _resume_notice_done, _rebindable_checkpointer

# None when checkpointing is inactive (the common, non-resume case).

checkpointer = inspect_ai.util.current_checkpointer()

# Only handoff runs re-register `compaction`, so only they need rebinding

if handoff and checkpointer is not None and _rebindable_checkpointer is None:

_rebindable_checkpointer = RebindableCheckpointer(checkpointer)

if _initial_messages is None:

captured = list(state.messages)

# `track` returns `captured` on a fresh run and the restored original

# input on resume; the reminder and compaction prefix derive from it.

# The callback closes over `tracked` (track's result) so a post-resume

# fire re-persists the original input, not the mid-run conversation.

tracked = captured

if checkpointer is not None:

tracked = checkpointer.track(

INITIAL_MESSAGES_CHECKPOINT_KEY,

lambda: tracked,

captured,

value_type=list[inspect_ai.model.ChatMessage],

)

_initial_messages = tracked

_input_messages = list(

msg

for msg in _initial_messages

if msg.role == "user" and msg.source == "input"

)

if _input_messages:

reminder_msg = inspect_ai.model.ContentText(

text="The following is a reminder of the original instructions you were given:\n\n"

)

# If no input messages in compacted context window, we append a reminder

# so the agent doesn't forget key info in the original instructions

_input_reminder_content = [reminder_msg] + list(

itertools.chain(

*[

[inspect_ai.model.ContentText(text=msg.content)]

if isinstance(msg.content, str)

else msg.content

for msg in _input_messages

]

)

)

if _tools is None:

_tools = tools

# Take over a handoff left by the handoff tools in the previous turn:

# the conversation restarts from the initial messages.

store = inspect_ai.util.store()

if handoff:

summary = store.get(metr_agents.tools.HANDOFF_SUMMARY_STORE_KEY)

if isinstance(summary, str):

store.set(metr_agents.tools.HANDOFF_SUMMARY_STORE_KEY, None)

handoff_count = (

store.get(metr_agents.tools.HANDOFF_COUNT_STORE_KEY, 0) + 1

)

store.set(metr_agents.tools.HANDOFF_COUNT_STORE_KEY, handoff_count)

inspect_ai.log.transcript().info(

{

"event": "handoff",

"handoff_count": handoff_count,

"summary_chars": len(summary),

},

source="metr_agents",

)

state.messages[:] = [

*_initial_messages,

inspect_ai.model.ChatMessageUser(

content=HANDOFF_NOTICE.format(summary=summary)

),

]

# Ensure that the timeout is reset to the default

store.set(

metr_agents.tools.TOOL_TIMEOUT_STORE_KEY,

metr_agents.tools.DEFAULT_TOOL_TIMEOUT,

)

# Compaction state describes the context we just cleared, so reset it

_compact = None

if _rebindable_checkpointer is not None:

# On resume, the checkpointed compaction state describes the

# pre-handoff conversation; replaying it would undo the wipe.

_rebindable_checkpointer.discard_checkpointed(

COMPACTION_CHECKPOINT_KEY

)

if _compact is None and compaction is not None:

_compact = build_compact_handler(

compaction,

initial_messages=_initial_messages,

tools=_tools,

checkpointer=_rebindable_checkpointer or checkpointer,

)

# optionally perform compaction on the input

if _compact is not None:

input_messages, c_message = await _compact.compact_input(state.messages)

# Summary and trim compaction, and OpenAI native compaction, preserve input

# messages, so we only append the input reminder for compaction methods that

# don't preserve input messages (currently only Anthropic native compaction)

# NB: need a fresh message id each time, or compaction mechanism will mark

# it as processed and then remove it every turn after second compaction

if (

not any(msg.role == "user" and msg.text for msg in input_messages)

and _input_reminder_content is not None

):

input_reminder = inspect_ai.model.ChatMessageUser(

id=shortuuid.uuid(), content=_input_reminder_content, source="input"

)

state.messages.append(input_reminder)

input_messages.append(input_reminder)

if c_message is not None:

state.messages.append(c_message)

else:

input_messages = state.messages

# On the first turn after a crash-resume, tell the model it was resumed:

# the work since the last checkpoint (including whatever crashed) was

# rolled back and is not shown. Once per process, appended after any

# compaction so it reaches this turn's generate call.

if not _resume_notice_done:

_resume_notice_done = True

if checkpointer is not None and checkpointer.attempt == "resume":

resume_notice = inspect_ai.model.ChatMessageUser(

id=shortuuid.uuid(), content=RESUME_NOTICE

)

state.messages.append(resume_notice)

if input_messages is not state.messages:

input_messages.append(resume_notice)

attempts = 0

while True:

# generate

output = await inspect_ai.model.get_model().generate(input_messages, tools)

# if it's a refusal see if we should retry

if output.stop_reason == "content_filter":

if retry_refusals is not None and attempts < retry_refusals:

attempts += 1

continue

# no retry, we are done

state.output = output

state.messages.append(state.output.message)

# update the compaction baseline with the actual input token

# count from the generate call (most accurate source of truth)

if _compact is not None:

await _compact.record_output(input_messages, output)

break

return state

return execute

def _resolve_tool_output_soft_limit(configured: int | None) -> int | None:

"""Resolve a configured max_tool_output into the soft limit our own tools

should truncate their output to.

Returns None when truncation should be disabled entirely (the user configured

a non-positive limit). An unset (None) limit falls back to the default.

"""

if configured is None:

return metr_agents.tools.DEFAULT_MAX_TOOL_OUTPUT

if configured <= 0:

return None

return configured

def _setup_tool_output_truncation() -> None:

"""Resolve this sample's tool-output soft limit, and take truncation over

from Inspect for the rest of the sample.

Inspect shares one GenerateConfig object across every sample of a

(task, model) run, and logs that same object as the eval's `plan.config` --

which `eval-retry` restores from and `eval_set` hashes into its task

identifier. So we must not mutate it. Instead we rebind the context variable

to a copy: context variables are re-bound per sample, so the copy is visible

to this sample only and the shared original is left untouched.

"""

config = inspect_ai.model._generate_config.active_generate_config()

configured = config.max_tool_output

if configured is None:

# The active generate config only carries eval/task-level values; a limit

# set on the model itself lives on the model's own config. Fall back to

# it, mirroring how Inspect merges the two (the eval-level value wins

# when set, otherwise the model-level one).

configured = inspect_ai.model.get_model().config.max_tool_output

soft_limit = _resolve_tool_output_soft_limit(configured)

if soft_limit is not None:

inspect_ai.util.store().set(

metr_agents.tools.MAX_TOOL_OUTPUT_STORE_KEY, soft_limit

)

# Our tools middle-truncate their own output to soft_limit (keeping the tail,

# which Inspect's truncator discards), so switch Inspect's pass off for this

# sample -- truncate_string_to_bytes treats max_bytes <= 0 as "no truncation".

inspect_ai.model._generate_config.set_active_generate_config(

config.model_copy(update={"max_tool_output": -1})

)

@inspect_ai.solver.solver

def react(

prompt: str

| dict[str, Any]

| inspect_ai.agent.AgentPrompt

| NotGiven

| None = NOT_GIVEN,

truncation: Literal["auto", "disabled"] | inspect_ai.agent.MessageFilter = "auto",

tools: metr_agents.tools.AgentToolSpec | None = None,

additional_tools: list[inspect_ai.tool.Tool] | None = None,

compaction: CompactionConfig | None = None,

submit: inspect_ai.agent.AgentSubmit | bool | None = None,

on_continue: str | inspect_ai.agent.AgentContinue | None = None,

limit_message_config: LimitMessageConfig

| dict[str, Any]

| NotGiven

| None = NOT_GIVEN,

handoff: bool = False,

implicit_tools: set[str] | None = None,

):

if isinstance(prompt, dict):

prompt = inspect_ai.agent.AgentPrompt(**prompt)

if on_continue is not None and not isinstance(limit_message_config, NotGiven):

raise ValueError(

"Cannot specify both 'on_continue' and 'limit_message_config'. "

+ "Use 'on_continue' to provide a custom callback, or "

+ "'limit_message_config' to customize the default limit messaging."

)

resolved_submit = (

submit

if submit is not None

else inspect_ai.agent.AgentSubmit(answer_delimiter=ANSWER_DELIMITER)

)

if on_continue is not None:

resolved_on_continue: str | inspect_ai.agent.AgentContinue = on_continue

elif isinstance(limit_message_config, NotGiven):

resolved_on_continue = limit_usage_message()

elif limit_message_config is None:

resolved_on_continue = DEFAULT_PROCEED_PROMPT

else:

limit_message_config = LimitMessageConfig.model_validate(limit_message_config)

resolved_on_continue = limit_usage_message(config=limit_message_config)

async def solve(

state: inspect_ai.solver.TaskState, generate: inspect_ai.solver.Generate

) -> inspect_ai.solver.TaskState:

_setup_tool_output_truncation()

optional_kwargs: OptionalReactKwargs = {}

if not isinstance(prompt, NotGiven):

optional_kwargs["prompt"] = prompt

tool_source = metr_agents.tools.TimeoutAwareDefaultToolSource(

existing_tools=state.tools,

tool_spec=tools,

additional_tools=additional_tools,

implicit_tools=implicit_tools,

)

return await inspect_ai.agent.as_solver(

inspect_ai.agent.react(

tools=[tool_source],

model=default_generator(compaction=compaction, handoff=handoff),

submit=resolved_submit,

on_continue=resolved_on_continue,

truncation=truncation,

**optional_kwargs,

)

)(state, generate)

return solve

def _constant_on_continue(message: str) -> inspect_ai.agent.AgentContinue:

"""Wrap a fixed continue message as a callback.

Agents built on `react(submit=False)` delegate to inspect's

`react_no_submit`, which rejects a plain `str` `on_continue` (with a submit

tool it would mean "the agent terminates when it makes no tool calls").

"""

async def on_continue(state: inspect_ai.agent.AgentState) -> str:

return message

return on_continue

async def _run_solver_and_append_stored_answer(

solver: inspect_ai.solver.Solver,

state: inspect_ai.solver.TaskState,

generate: inspect_ai.solver.Generate,

*store_keys: str,

) -> inspect_ai.solver.TaskState:

"""Run `solver`, then append the answer left in the store to the completion.

`store_keys` are tried in order and the first one that holds a value wins,

so precedence is fixed by the caller rather than by which tool happened to

write last.

The append happens even when the solver raises or a limit fires: although we

don't return `state` in that case, it is the sample's state and is mutated

in place, so the change persists.

"""

try:

state = await solver(state, generate)

finally:

store = inspect_ai.util.store()

answer = next(

(value for key in store_keys if (value := store.get(key)) is not None),

"",

)

state.output.completion = f"{state.output.completion}{ANSWER_DELIMITER}{answer}"

return state

CHECKPOINT_SUBMIT_TOOL_NAME = "metr_agents/checkpoint_submit"

CHECKPOINT_PROCEED_PROMPT = (

"\nPlease proceed to the next step using your best judgement.\n"

)

CHECKPOINT_LIMIT_MESSAGE_CONFIG = LimitMessageConfig(

defaults=LimitWarningDefaults(

near=LimitWarningConfig(

threshold=0.80,

message="You are nearing the limit. Prepare to wrap up.",

),

very_near=LimitWarningConfig(

threshold=0.95,

message="You are almost at the limit. Wrap up as soon as possible.",

),

),

)

@inspect_ai.solver.solver

def react_with_checkpoint_submit(

prompt: str

| dict[str, Any]

| inspect_ai.agent.AgentPrompt

| NotGiven

| None = NOT_GIVEN,

truncation: Literal["auto", "disabled"] | inspect_ai.agent.MessageFilter = "auto",

tools: metr_agents.tools.AgentToolSpec | None = None,

compaction: CompactionConfig | None = None,

limit_message_config: LimitMessageConfig

| dict[str, Any]

| NotGiven

| None = NOT_GIVEN,

):

if isinstance(limit_message_config, NotGiven):

resolved_on_continue = limit_usage_message(

proceed_prompt=CHECKPOINT_PROCEED_PROMPT,

config=CHECKPOINT_LIMIT_MESSAGE_CONFIG,

)

elif limit_message_config is None:

resolved_on_continue = _constant_on_continue(CHECKPOINT_PROCEED_PROMPT)

else:

resolved_on_continue = limit_usage_message(

proceed_prompt=CHECKPOINT_PROCEED_PROMPT,

config=LimitMessageConfig.model_validate(limit_message_config),

)

solver = react(

prompt=prompt,

truncation=truncation,

tools=tools,

compaction=compaction,

submit=False,

on_continue=resolved_on_continue,

additional_tools=[metr_agents.tools.checkpoint_submit()],

implicit_tools={CHECKPOINT_SUBMIT_TOOL_NAME},

# This tool stands in for inspect's submit tool, which a tool spec never

# has to configure; requiring it would be a trap for task authors.

)

async def solve(

state: inspect_ai.solver.TaskState, generate: inspect_ai.solver.Generate

) -> inspect_ai.solver.TaskState:

limits = inspect_ai.util.sample_limits()

if (

limits.token.limit is None

and limits.time.limit is None

and limits.working.limit is None

and limits.cost.limit is None

):

raise ValueError(

"No token_limit, time_limit, working_limit, or cost_limit is set. "

+ "You must set at least one limit to use react_with_checkpoint_submit."

)

return await _run_solver_and_append_stored_answer(

solver, state, generate, metr_agents.tools.CHECKPOINT_STORE_KEY

)

return solve

@inspect_ai.solver.solver

def react_with_gated_submit(

prompt: str

| dict[str, Any]

| inspect_ai.agent.AgentPrompt

| NotGiven

| None = NOT_GIVEN,

truncation: Literal["auto", "disabled"] | inspect_ai.agent.MessageFilter = "auto",

tools: metr_agents.tools.AgentToolSpec | None = None,

compaction: CompactionConfig | None = None,

gated_submit_token_fraction: float = 0.5,

early_submit_message: str = DEFAULT_EARLY_SUBMIT_MESSAGE,

proceed_prompt: str | None = DEFAULT_PROCEED_PROMPT,

limit_message_config: LimitMessageConfig

| dict[str, Any]

| NotGiven

| None = NOT_GIVEN,

):

if not 0.0 <= gated_submit_token_fraction < 1.0:

raise ValueError(

f"token_fraction must be in [0.0, 1.0), got {gated_submit_token_fraction}"

)

if isinstance(limit_message_config, NotGiven):

resolved_on_continue: str | inspect_ai.agent.AgentContinue = (

limit_usage_message(proceed_prompt=proceed_prompt)

)

elif limit_message_config is None:

resolved_on_continue = proceed_prompt or ""

else:

resolved_on_continue = limit_usage_message(

proceed_prompt=proceed_prompt,

config=LimitMessageConfig.model_validate(limit_message_config),

)

_solver = react(

prompt=prompt,

truncation=truncation,

tools=tools,

compaction=compaction,

on_continue=resolved_on_continue,

submit=inspect_ai.agent.AgentSubmit(

name="submit",

tool=metr_agents.tools.gated_submit(

token_fraction=gated_submit_token_fraction,

early_submit_message=early_submit_message,

),

answer_delimiter=ANSWER_DELIMITER,

),

)

async def solve(

state: inspect_ai.solver.TaskState, generate: inspect_ai.solver.Generate

) -> inspect_ai.solver.TaskState:

if state.token_limit is None:

raise ValueError(

"gated_submit_token_fraction is set but token_limit is None. "

+ "You must set a token limit to use react_with_gated_submit."

)

return await _solver(state, generate)

return solve

HANDOFF_NOTICE = """You are taking over this task from another agent. The environment is exactly as the other agent left it.

The previous agent left this summary of its work:

<handoff_summary>

{summary}

</handoff_summary>

Do not take the summary's assessment of progress or completeness at face value. Study the current state of the work yourself, read the original task instructions, find any limitations and ways that the task might not be fully completed, then improve it.

"""

SUBMIT_TOOL_NAME = "metr_agents/submit"

HANDOFF_TOOL_NAME = "metr_agents/handoff"

HANDOFF_SUBMIT_PROMPT = f"""

When you have completed the task and have an answer, call the submit() tool to report it. Along with your answer, provide a summary of at most {metr_agents.tools.MAX_HANDOFF_SUMMARY_CHARS:,} characters describing the work you have done: what you learned, the current state, and anything left unfinished. If there is remaining budget, a fresh copy of you with no memory of this conversation will continue the task in this same sandbox, starting from that summary.

"""

HANDOFF_OR_EARLY_SUBMIT_PROMPT = f"""

When you wish to hand off your task, call the handoff() tool. It records your current best answer. Along with your answer, provide a summary of at most {metr_agents.tools.MAX_HANDOFF_SUMMARY_CHARS:,} characters describing the work you have done: what you learned, the current state, and what remains to improve. If there is remaining budget, a fresh copy of you with no memory of this conversation continues working in this same sandbox, starting from that summary. You should have a very high bar for describing the work as complete. Focus the summary on limitations of the current state and concrete areas of improvement, so the next instance knows where to direct its effort.

Only call the submit() tool if you are absolutely confident that the task is fully completed and additional work cannot improve the deliverable. Submitting ends the run permanently. No further work on the task is possible afterwards, and any remaining budget is forfeited. When in doubt about whether to submit or hand off, err on the side of handing off.

"""

# In `allow_early_submit` mode the default proceed prompt and limit warnings

# would tell the model to call `submit()`, which ends the run and forfeits the

# remaining budget. Handing off records the same answer and keeps working, so

# that is what the model should be urged towards.

HANDOFF_PROCEED_PROMPT = (

"\nPlease proceed to the next step using your best judgement. If you have\n"

"taken this line of work as far as you usefully can, call the `handoff()`\n"

"tool so a fresh instance can continue from your summary.\n"

)

HANDOFF_LIMIT_MESSAGE_CONFIG = LimitMessageConfig(

defaults=LimitWarningDefaults(

near=LimitWarningConfig(

threshold=0.80,

message="You are nearing the limit. Prepare for the work to be submitted.",

),

very_near=LimitWarningConfig(

threshold=0.95,

message=(

"You are almost at the limit. Call `submit()` as soon as possible."

),

),

)

)

def _validate_handoff_limits() -> None:

"""Check the sample has a limit that can actually end a handoff run.

Handing off never terminates, and even with `allow_early_submit` the

agent may hand off forever without submitting, so only a limit guarantees

the run ends. `message_limit` cannot help: a handoff clears the

conversation and inspect checks the limit against the current conversation

length, so every handoff resets the count.

"""

limits = inspect_ai.util.sample_limits()

if limits.message.limit is not None:

raise ValueError(

"message_limit cannot be used with react_with_handoff_submit. Each "

+ "handoff clears the conversation, which resets the message count, "

+ "so a message limit can never end the run. Use token_limit, "

+ "time_limit, working_limit, or cost_limit instead."

)

if (

limits.token.limit is None

and limits.time.limit is None

and limits.working.limit is None

and limits.cost.limit is None

):

raise ValueError(

"No token_limit, time_limit, working_limit, or cost_limit is set. "

+ "You must set at least one limit to use react_with_handoff_submit."

)

@inspect_ai.solver.solver

def react_with_handoff_submit(

prompt: str

| dict[str, Any]

| inspect_ai.agent.AgentPrompt

| NotGiven

| None = NOT_GIVEN,

truncation: Literal["auto", "disabled"] | inspect_ai.agent.MessageFilter = "auto",

tools: metr_agents.tools.AgentToolSpec | None = None,

compaction: CompactionConfig | None = None,

limit_message_config: LimitMessageConfig

| dict[str, Any]

| NotGiven

| None = NOT_GIVEN,

allow_early_submit: bool = False,

):

if allow_early_submit and not isinstance(prompt, NotGiven):

raise ValueError(

"A custom prompt cannot be combined with allow_early_submit=True. "

+ "In that mode `submit()` ends the run permanently and `handoff()` "

+ "continues it, and the agent will not behave correctly without the "

+ "default prompt that explains how to use these tools."

)

if isinstance(limit_message_config, NotGiven):

resolved_on_continue = limit_usage_message(

proceed_prompt=HANDOFF_PROCEED_PROMPT

if allow_early_submit

else DEFAULT_PROCEED_PROMPT,

config=HANDOFF_LIMIT_MESSAGE_CONFIG if allow_early_submit else None,

)

elif limit_message_config is None:

resolved_on_continue = _constant_on_continue(

HANDOFF_PROCEED_PROMPT if allow_early_submit else DEFAULT_PROCEED_PROMPT

)

else:

config = LimitMessageConfig.model_validate(limit_message_config)

if allow_early_submit and "defaults" not in config.model_fields_set:

config = config.model_copy(

update={"defaults": HANDOFF_LIMIT_MESSAGE_CONFIG.defaults}

)

resolved_on_continue = limit_usage_message(

proceed_prompt=HANDOFF_PROCEED_PROMPT

if allow_early_submit

else DEFAULT_PROCEED_PROMPT,

config=config,

)