fix(session): respect configured thinking in patch merge and repair budget-exhausted extraction

[akushonkamen]

Description

PatchMergePolicyOptimizer hard-codes a per-call thinking=True for experiences/trajectories merges (openviking/session/train/components/policy_optimizer.py:186 on main), and the volcengine backend lets an explicit per-call value override its configured thinking policy (openviking/models/vlm/backends/volcengine_vlm.py:246: effective_thinking = self.thinking if thinking is None else thinking; config default False at openviking/models/vlm/base.py:79). A deployment with thinking=False therefore still had reasoning forced on during patch merge.

The reasoning consumed the entire completion budget: the output was truncated mid-string (finish_reason=length), but the extract loop's single format retry classified it as an ordinary parse_error — it did not look at finish_reason/completion_tokens, did not append the failed assistant program, and reused the same max_tokens. The retry reproduced the identical truncation and the session commit failed terminally after ~2h.

Root cause

1. Per-call thinking=True in PatchMergePolicyOptimizer overrides a deployment's configured thinking=False (per-call value wins over config in the volcengine backend).
2. The extract loop's single format retry is budget-blind: truncation (finish_reason == "length" or completion_tokens >= max_tokens) is treated the same as a syntax mistake, the truncated assistant program is not shown to the model, and the output cap is not adjusted — so a budget-exhausted generation can never be repaired.

Fix

• policy_optimizer: new thinking: bool | None = None field. The default (None) defers to the VLM backend's configured thinking policy; explicit True/False still force the mode for callers that want it.
• extract_loop: signature changed to thinking: Optional[bool] = None and passed through verbatim (so None is not blocked); existing callers passing explicit booleans are unchanged.
• On parse failure, the response is classified as budget_exhausted when finish_reason == "length" or completion_tokens >= the effective output cap. In that case the single retry: appends the failed assistant program so the model repairs its own truncated output, renders protocol-level budget guidance via the new render_budget_retry (default in extraction_output_protocol/base.py, with a python-protocol SDK-specific override), and raises the cap once, bounded to min(2x, 65536) and never lowered. Ordinary format mistakes keep the previous repair behavior; malformed Python programs are still rejected.

Human Involvement

• A human participated in the implementation or review loop
• This PR was generated entirely by AI agents without human participation in the loop

Related Issue

Fixes #6096

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Refactoring (no functional changes)
• Performance improvement
• Test update

Changes Made

• openviking/session/train/components/policy_optimizer.py: add thinking field (bool | None, default None) and pass it to the extract loop instead of forcing thinking=self.memory_type in {"trajectories", "experiences"}.
• openviking/session/memory/extract_loop.py: accept Optional[bool] thinking and forward it verbatim; classify budget-truncated responses as budget_exhausted; on budget exhaustion make the single retry append the failed assistant program, render render_budget_retry guidance, and raise the output cap once (bounded min(2x, 65536)).
• openviking/session/memory/extraction_output_protocol/base.py: add default render_budget_retry.
• openviking/session/memory/extraction_output_protocol/python_protocol.py: python-protocol-specific render_budget_retry override.
• tests/session/memory/test_extract_loop_budget_recovery.py: new (7 tests) — budget classification and single-retry recovery behavior.
• tests/session/train/test_train_components.py: new wiring tests (3) — optimizer thinking defaults to config passthrough; explicit values are forwarded.

Testing

• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have tested this on the following platforms:
  • Linux
  • macOS
  • Windows

All commands below ran with OPENVIKING_CONFIG_FILE=/tmp/ovtest/ov.conf uv run --no-sync python -m pytest … --override-ini addopts="".

Red (before the fix, re-verified on a clean 89fbcb9 worktree with the new test files applied): 7 of the 10 new tests fail — tests/session/memory/test_extract_loop_budget_recovery.py 4 failed (3 passed: they pin the ordinary parse-error behavior that already works on main); tests/session/train/test_train_components.py -k thinking 3 failed.

Green (after the fix, rebased onto current origin/main = 9baf83f, this push's HEAD ad1bbae):

• pytest tests/session/memory/test_extract_loop_budget_recovery.py → 7 passed
• pytest tests/session/train/test_train_components.py -k thinking → 3 passed (23 deselected)
• Related subset (test_memory_react.py test_extraction_output_protocol.py test_streaming_memory_updater.py test_consolidation_compile.py test_train_components.py test_extract_loop_budget_recovery.py) → 284 passed, 2 failed; the 2 failures are test_skill_policy_creation_preserves_lock_ownership[standalone|outer-tree-lock], the ragfs_python environment noise already documented on the issue.
• Full regression pytest tests/session/ -q: clean origin/main (9baf83f) baseline → 1244 passed, 36 failed, 5 skipped, 118 errors; with this fix → 1254 passed, 36 failed, 5 skipped, 118 errors. The FAILED/ERROR lists (154 lines each) are byte-identical after diff, i.e. zero new failures; the +10 passed are exactly the new tests.

Checklist

• My code follows the project's coding style
• I have performed a self-review of my code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• Any dependent changes have been merged and published

Additional Notes

• AI disclosure: this PR (analysis, patch, tests) was produced entirely by AI agents without human participation in the loop, per the Human Involvement section above.
• The 36 failed / 118 errors in tests/session/ are pre-existing on clean origin/main and untouched by this PR (identical failure lists; includes the 2 lock_ownership failures caused by the ragfs_python environment, as noted on the issue).

[cocolord]

Thanks for picking this up. The thinking part looks right to me: keeping None / False / True and letting None inherit the VLM config is exactly what the issue asks for.

The truncation half doesn't hold on the actual incident path, though. I checked it against ad1bbae:

1. length never reaches _response_hit_output_budget on patch merge. PatchMergeContextProvider.get_tools() returns [], so _call_llm sends tools=None. The Volcengine and LiteLLM backends then return a plain str and drop finish_reason and usage. _response_hit_output_budget returns False for str. With the real VolcEngineVLM returning finish_reason="length" and completion_tokens=32768, the failure is still classified as parse_error. The new tests pass because ScriptedVLM returns a VLMResponse whatever tools is.
2. A truncated program that still parses is accepted. The budget check only runs after a parse failure. A program cut off at a statement boundary is valid Python, so finish_reason="length" with a parseable body is returned as successful operations. The trailing writes and deletes are silently lost.
3. completion_tokens >= cap is treated as truncation. A finish_reason="stop" response that lands exactly on the cap is classified as budget_exhausted. The issue explicitly avoids inferring finish_reason from token counters.

Repro (all three fail on ad1bbae)

from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest

from openviking.models.vlm.backends.volcengine_vlm import VolcEngineVLM
from openviking.models.vlm.base import VLMResponse
from openviking.session.memory.extract_loop import ExtractLoop


def _loop(vlm, parse_result):
    loop = ExtractLoop(
        vlm=vlm, viking_fs=MagicMock(), context_provider=MagicMock(), max_output_tokens=32768
    )
    loop._tool_schemas = []  # PatchMergeContextProvider.get_tools() == []
    loop._output_protocol = MagicMock()
    loop._output_protocol.parse.return_value = parse_result
    loop._output_context = MagicMock()
    return loop


def _fake_vlm(response):
    vlm = MagicMock(model="test-model")

    async def get_completion_async(**_kwargs):
        return response

    vlm.get_completion_async = get_completion_async
    return vlm


@pytest.mark.asyncio
async def test_length_is_detected_through_real_backend_on_patch_merge_path():
    raw = SimpleNamespace(
        choices=[
            SimpleNamespace(
                message=SimpleNamespace(content='sdk.write("a", """unterminated', tool_calls=None),
                finish_reason="length",
            )
        ],
        usage=SimpleNamespace(
            prompt_tokens=20214,
            completion_tokens=32768,
            total_tokens=52982,
            prompt_tokens_details=None,
            completion_tokens_details=None,
        ),
    )

    async def create(**_kwargs):
        return raw

    vlm = VolcEngineVLM({"provider": "volcengine", "model": "test-model"})
    vlm.get_async_client = lambda: SimpleNamespace(
        chat=SimpleNamespace(completions=SimpleNamespace(create=create))
    )
    loop = _loop(vlm, (None, "unterminated triple-quoted string literal"))

    await loop._call_llm([])

    assert loop._last_llm_failure_kind == "budget_exhausted"  # actual: parse_error


@pytest.mark.asyncio
async def test_length_response_that_still_parses_is_not_accepted():
    operations = object()
    loop = _loop(
        _fake_vlm(VLMResponse(content="sdk.write('a', 'x')", finish_reason="length")),
        (operations, None),
    )

    _tool_calls, result = await loop._call_llm([])

    assert result is not operations  # actual: truncated program accepted


@pytest.mark.asyncio
async def test_stop_at_cap_is_not_truncation():
    loop = _loop(
        _fake_vlm(
            VLMResponse(content="bad", finish_reason="stop", usage={"completion_tokens": 32768})
        ),
        (None, "invalid syntax"),
    )

    await loop._call_llm([])

    assert loop._last_llm_failure_kind != "budget_exhausted"  # actual: budget_exhausted

Suggested direction:

• Have _call_llm pass tools=[] instead of None when tools are disabled. Align Volcengine and LiteLLM with the OpenAI backend's existing tools is not None convention from #3535, so a no-tool call still returns a VLMResponse with finish_reason and usage. When tools is empty, those backends must still run _clean_response on the content. Otherwise inline <think>...</think> output from OpenAI-compatible or LiteLLM models stops parsing.
• Check for an explicit finish_reason == "length" before parsing, and fail with a dedicated, actionable error (cap, usage, preview). Don't route it through the format retry. Don't use completion_tokens as a truncation signal on its own.
• I'd drop the cap doubling and the replay of the truncated output from this PR. The issue doesn't assume a larger cap or a repeated prompt will help. Replaying ~32k tokens of truncated output also grows the prompt. Once thinking is no longer forced, failing fast is enough here, and any bounded recovery can be designed separately.

The main things worth looking at: the real-backend test for the tools=None / str path, and the parseable-but-truncated case.

[akushonkamen]

Thanks for the thorough review and the repro — these are valid findings, and I agree the current truncation handling doesn't hold on the real patch-merge path. I'll rework it along your suggested direction: pass tools=[] instead of None in _call_llm and align the Volcengine/LiteLLM backends with the OpenAI backend's tools is not None convention (keeping _clean_response on empty-tools content), check an explicit finish_reason == "length" before parsing with a dedicated actionable error outside the format-retry path, and drop the completion_tokens-based heuristic. I'll also remove the cap doubling and the truncated-output replay as you suggest. Will push a revision.

[akushonkamen]

Pushed 4f18d6c reworking the truncation handling along your suggested direction. Point by point:

1. tools=[] instead of None + backend convention. _call_llm now passes tools=[] when tools are disabled (openviking/session/memory/extract_loop.py:1210). The Volcengine and LiteLLM backends' return paths now follow the OpenAI backend's tools is not None convention (openviking/models/vlm/backends/volcengine_vlm.py, openviking/models/vlm/backends/litellm_vlm.py): tools=[] is still not sent on the wire, but the response is built as a VLMResponse carrying finish_reason and usage. When tools is empty, those backends still run _clean_response on the content, so inline <think>...</think> output keeps parsing. Regression tests: test_empty_tools_response_from_real_volcengine_backend (real VolcEngineVLM with a mocked client, tools=[], <think>-wrapped content → VLMResponse with cleaned content and preserved finish_reason/usage) and test_length_through_real_volcengine_backend_is_budget_exhausted (your main repro: finish_reason="length" via the real backend on the patch-merge path is now classified budget_exhausted, not parse_error).

2. Explicit finish_reason == "length" check before parsing. _call_llm now checks finish_reason == "length" right after the response is received and before parse() runs, and records a dedicated actionable error naming the finish reason, the output cap, completion_tokens and a preview, plus the fix (vlm.max_tokens / smaller context) (extract_loop.py _budget_exhausted_error). run() treats budget_exhausted as terminal: it raises immediately and never routes through the format retry (test_length_finish_reason_is_terminal_even_when_program_parses — a length-truncated program that would parse cleanly is rejected, so trailing writes/deletes are not silently lost). completion_tokens is no longer used as a truncation signal anywhere: test_stop_at_cap_is_not_truncation covers your third repro (finish_reason="stop" landing exactly on the cap is an ordinary format error with a normal retry).

3. Cap doubling and truncated-output replay dropped. Removed _raise_output_budget_for_retry, _EXTRACTION_MAX_OUTPUT_TOKENS_CEILING, the assistant-message replay of the truncated program in _add_format_error_message, and render_budget_retry from both output protocols. The truncated output is never appended to the prompt; failing fast is the whole recovery now.

Test notes: the rewritten tests/session/memory/test_extract_loop_budget_recovery.py passes 10/10 locally (.venv/bin/python -m pytest tests/session/memory/test_extract_loop_budget_recovery.py). One harness fix was needed: _run_loop previously returned the run() coroutine outside the patch(...) context, so the config patches had already expired when the loop executed — every test in this file failed on the previous head for that reason; it now awaits inside the context. The wider tests/session/ tree has ~318 pre-existing failures in this environment (config-dependent tests, identical failures on the base commit — verified by stash/baseline runs); my change moves that count from 325 failed/778 passed to 318 failed/788 passed with no new failures, but I could not get those suites green locally, so the wider suite remains unvalidated beyond that comparison.

[J-StrawHat]

Thanks for the update.

I reviewed 4f18d6c8 and found two remaining issues:

1. Passing tools=[] bypasses content cleaning in OpenAIVLM.get_completion_async, leaving inline <think>...</think> content intact. An otherwise valid SDK program can therefore fail with parse_error.

2. Responses containing tool calls return before the finish_reason == "length" check. A truncated response can still execute a tool and trigger another model call instead of failing immediately.

I've addressed both in a small follow-up commit: 20ce905. It restores content cleaning for the empty-tools async text-completion path and moves the existing truncation check ahead of tool-call handling.