Skip to content

fix(reflect): pair each expanded memory_id with its own memory - #2759

Merged
benfrank241 merged 1 commit into
vectorize-io:mainfrom
ebarkhordar:fix/reflect-expand-memory-id-pairing
Jul 17, 2026
Merged

fix(reflect): pair each expanded memory_id with its own memory#2759
benfrank241 merged 1 commit into
vectorize-io:mainfrom
ebarkhordar:fix/reflect-expand-memory-id-pairing

Conversation

@ebarkhordar

Copy link
Copy Markdown
Contributor

Root cause

tool_expand validates memory_ids into a second, shorter list and then zips the two back together:

valid_uuids: list[uuid.UUID] = []
errors: dict[str, str] = {}
for mid in memory_ids:
    try:
        valid_uuids.append(uuid.UUID(mid))     # only the ids that PARSE land here
    except ValueError:
        errors[mid] = f"Invalid memory_id format: {mid}"
...
for mid, mem_uuid in zip(memory_ids, valid_uuids):   # tools.py:398

valid_uuids is compacted, so valid_uuids[i] lines up with memory_ids[i] only while every preceding id parsed. One invalid id shifts every later pair by one, and zip then truncates to the shorter list, dropping the tail.

Two things follow, and the first is the one that worries me:

  1. A memory is returned stamped with a different memory's id. item["memory_id"] and item["memory"]["id"] disagree inside the same response item, so nothing raises and nothing logs: the reflect agent is handed B's text under A's id and cites it as A.
  2. The last requested id gets no entry at all, not even an error entry, so the errors branch that exists to report a bad id silently does not fire for it.

Invariant

tool_expand returns one entry per requested memory_id, in request order, and each entry's memory payload is the row whose id equals that entry's own memory_id.

Measured on HEAD (9676fc1)

--- CASE 1: ["not-a-uuid", A, B]
    count returned: 2                          <- 3 requested
    entry memory_id=not-a-uuid -> error='Invalid memory_id format: not-a-uuid'
    entry memory_id=aaaaaaaa-...-aaaaaaaaaaaa
          memory.id=bbbbbbbb-...-bbbbbbbbbbbb  <-- MIS-ATTRIBUTION
          text='MEMORY-B: dogs are loyal animals'
    MISSING (requested, no entry at all): ['bbbbbbbb-...-bbbbbbbbbbbb']

--- CASE 2: [A, "bad"]
    count returned: 1
    MISSING (requested, no entry at all): ['bad']

--- CASE 3 (control): [A, B] -> both correct

Case 3 is the control: with every id valid the pairing is correct, which is why the compaction is the trigger rather than the lookup.

Reachability

_execute_tool (agent.py:1420-1425) reads args.get("memory_ids", []), checks only that it is non-empty, and passes the raw model-supplied array to expand_fn; tools_schema.py:121-125 declares memory_ids as a free-form array of string. So one hallucinated or truncated id among valid ones is enough, and the errors dict plus the if mid in errors branch already in the function are the existing acknowledgement that such input arrives.

The fix

Key each id to its own UUID and iterate memory_ids directly, so an invalid id can only affect its own entry. Nine lines added and six removed in tools.py, no signature or API change.

One incidental, measured rather than assumed: valid_uuids is now built from the dict's values, so a repeated id is sent to WHERE id = ANY($1) once instead of once per occurrence. The returned results are unchanged (a duplicate id still gets one entry per occurrence, verified both before and after).

How I verified

Clean python:3.12-slim container, package installed non-editable, HEAD 9676fc1. Provenance was pinned by sha256 identity between the installed module and the checkout, both sides of the differential:

  • unfixed: 6f05a294f95600e6ad80573b949b78b90b070f93e35fa564589a254b185684b4
  • fixed: 82d80e5f771d4399c918768b752be037b8257f45030827bded716eca502b89c5

The two regression tests fail on main and pass on the branch (run both ways in the container, against the printed sha256 above):

main:   2 failed    (assert 1 == 2; assert 2 == 3)
branch: 2 passed
tests/test_reflect_tools.py in full, on branch: 16 passed

The tests use this file's existing fake-connection idiom, so they need no database and no LLM. The connection is the only thing faked, and it supplies rows; the pairing under test happens entirely inside tool_expand and runs for real.

ruff check and ruff format --check pass on both changed files. ty check hindsight_api reports 128 diagnostics on my branch and the same 128 on main, none in reflect/tools.py.

Each of the two behaviors named above has its own test: test_tool_expand_pairs_each_memory_id_with_its_own_memory (case 1) and test_tool_expand_reports_a_trailing_invalid_memory_id (case 2).

What I did not verify

  • Not run against a real Postgres. The connection is faked, matching the existing tests in this file.
  • I did not audit the other 27 zip() call sites in the package for the same shape; this PR is only tool_expand.
  • I did not measure how often a model actually emits a malformed id here, so I cannot say how often this fires in production. The argument for the fix is that the code already has an error path for that input and mis-pairs instead of taking it.

git log -S 'zip(memory_ids' on an unshallowed clone returns exactly one commit, 4f28338 (#132), which introduced the tool. This is an original defect rather than a regression, and there is no earlier attempt or revert to answer.

Disclosure: I am an AI agent (Claude) working on Ehsan Barkhordar's behalf. He is accountable for this PR. Every number above came from a run in the container described, not from inspection.

tool_expand zipped memory_ids against valid_uuids, which only collects the
ids that parsed as UUIDs. One invalid id shifts every later pair by one, so
a memory comes back stamped with a different memory's id, and zip truncates
the tail so the last requested id gets no entry at all.

Key each id to its own UUID and iterate memory_ids directly, so an invalid
id can only affect its own entry.
@benfrank241
benfrank241 merged commit 1fe43ec into vectorize-io:main Jul 17, 2026
87 checks passed
@ebarkhordar
ebarkhordar deleted the fix/reflect-expand-memory-id-pairing branch July 17, 2026 16:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants