chore: inject_state_manager should accept max_extend_iters and extend_iters_per_candidate as parameters

[fatihbugrakdogan]

Problem

inject_state_manager reads max_extend_iters and extend_iters_per_candidate exclusively from environment variables. There is no way to pass these values as function parameters.

In [kai/state/integration.py L119-L132]:

iters_per_candidate = int(
    os.environ.get(
        "KAI_EXTEND_ITERS_PER_CANDIDATE",
        _DEFAULT_EXTEND_ITERS_PER_CANDIDATE,
    )
)
# ...
extras["max_iterations_limit"] = config.max_iterations + int(
    os.environ.get("KAI_MAX_EXTEND_ITERS", _DEFAULT_MAX_EXTEND_ITERS)
)

The function signature has no parameters for these values (L39-L47):

  def inject_state_manager(
      config: RecursiveAgentConfig,
      state_manager: StateManager,
      run_id: str,
      result_processors: dict[str, ResultProcessor] | None = None,
      *,
      save_rollouts: bool = False,
      rollout_agents: set[str] | None = None,
      _depth: int = 0,
  ) -> RecursiveAgentConfig:

Even though RecursiveAgentConfig already has max_iterations_limit and on_extend as fields (ra/agents/config.py L38-L39):

  on_extend: Callable[[int], int | None] | None = None
  max_iterations_limit: int | None = None

setting them on the config before calling inject_state_manager has no effect, because the function unconditionally overwrites them via

replace(config, **extras) at L135-L140.

Why this matters

In kai-executor, these values are now stored per-tier in MongoDB (maxExtendIters, extendItersPerCandidate on the tier document). The executor loads them into TierConfig but is forced to inject them through os.environ before calling inject_state_manager:

  os.environ["KAI_MAX_EXTEND_ITERS"] = str(self.tier_config.max_extend_iters)
  os.environ["KAI_EXTEND_ITERS_PER_CANDIDATE"] = str(self.tier_config.extend_iters_per_candidate)
  injected_config = inject_state_manager(injected_config, self.state_manager, run_id)

This is fragile — environment variables are global mutable state, not scoped to the call.

Proposed solution

Add optional keyword arguments to inject_state_manager, falling back to env vars then defaults when not provided:

  def inject_state_manager(
      config: RecursiveAgentConfig,
      state_manager: StateManager,
      run_id: str,
      result_processors: dict[str, ResultProcessor] | None = None,
      *,
      save_rollouts: bool = False,
      rollout_agents: set[str] | None = None,
      max_extend_iters: int | None = None,
      extend_iters_per_candidate: int | None = None,
      _depth: int = 0,
  ) -> RecursiveAgentConfig:

This way callers can pass values directly while existing env var behavior is preserved for backward compatibility.