infer.py: get_query() on full sequence may match <search> block from previous turns

[azure-z77]

Description

In infer.py, line 118, the get_query() function is called on the full decoded sequence (prompt + all history + new tokens), but it should only search within the newly generated tokens. This can cause incorrect behavior in multi-turn search scenarios.

Problematic Code

https://github.com/PeterGriffinJin/Search-R1/blob/main/search_r1/infer.py

# Line 116: correctly decodes only the newly generated tokens
output_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)

# Line 118: decodes the ENTIRE sequence (prompt + history + new tokens)
tmp_query = get_query(tokenizer.decode(outputs[0], skip_special_tokens=True))

The get_query() function uses re.findall and takes matches[-1] (the last match):

def get_query(text):
    pattern = re.compile(r"<search>(.*?)</search>", re.DOTALL)
    matches = pattern.findall(text)
    if matches:
        return matches[-1]  # returns the LAST <search> block
    else:
        return None

Root Cause

In the while True loop, each turn's output and search results are appended back to prompt (line 126). So outputs[0] after decoding contains all previous turns' <search>...</search> blocks.

If the model hits max_new_tokens (1024) while still in the <think> phase — without producing a new <search> block — get_query() on the full sequence will match a stale block from a previous turn, causing:

• The same old query to be searched again
• Duplicate search results to be appended to the prompt, polluting subsequent generation

Proposed Fix

Use output_text (the already-computed new-tokens-only string) instead of re-decoding the full sequence:

# Before (line 118):
tmp_query = get_query(tokenizer.decode(outputs[0], skip_special_tokens=True))

# After:
tmp_query = get_query(output_text)

output_text is already defined on line 116 and contains only the newly generated tokens, so it will only match a block produced in the current turn.

Impact

This is an edge case that occurs when max_new_tokens is reached mid-generation without a stop trigger. While not common, it can lead to silent incorrect behavior (stale query reuse) that is hard to notice.