feat(logits): add safety for zero copy access

[lloyal-research]

No description provided.

[greptile-apps]

Greptile Overview

Greptile Summary

This PR implements a memory safety mechanism for zero-copy logits access using a "Memoized Step-Scoped Views with Explicit Revocation" pattern. The changes prevent use-after-invalidation bugs when JavaScript code holds references to logits buffers that become stale after decode() or encode() calls.

Key Changes:

• Added invalidateLogits() method that detaches ArrayBuffers before operations that invalidate logits (decode, encode, dispose)
• Implemented memoization so multiple getLogits() calls in the same step return the same buffer
• Added withLogits() helper in JavaScript that enforces synchronous-only usage with runtime checking
• Created new lloyal::logits::get() wrapper with null checking and descriptive error messages
• Comprehensive test coverage for buffer revocation and memoization behavior

Issue Found:

• src/SessionContext.cpp:699 - Reference count is set to 1 (strong reference) but comment says "weak reference". Should use 0 to avoid preventing garbage collection.

Confidence Score: 3/5

• Safe to merge after fixing reference count issue in SessionContext.cpp:699
• The PR implements a solid memory safety pattern with comprehensive tests. However, there's a critical bug where the reference count is 1 instead of 0, creating a strong reference instead of weak reference. This will prevent the buffer from being garbage collected, causing a memory leak. Once this single-line fix is applied, the implementation is sound.
• src/SessionContext.cpp requires immediate attention - reference count must be changed from 1 to 0 on line 699

Important Files Changed

File Analysis

Filename	Score	Overview
src/SessionContext.cpp	3/5	added logits buffer management with memoization and revocation. critical issue: reference count should be 0 for weak reference, not 1
src/SessionContext.hpp	5/5	added logits buffer tracking fields and invalidation method. clean implementation with good documentation
lib/index.js	5/5	added withLogits helper with runtime borrow checker pattern. properly detects and prevents async usage
lib/index.d.ts	5/5	added TypeScript types for withLogits helper with comprehensive documentation and examples
vendor/liblloyal/include/lloyal/logits.hpp	5/5	added safe logits wrapper with null checking and clear lifetime documentation. good defensive programming

Sequence Diagram

sequenceDiagram
    participant JS as JavaScript Code
    participant Ctx as SessionContext
    participant Buf as ArrayBuffer
    participant CPP as llama.cpp

    Note over JS,CPP: Initial Decode
    JS->>Ctx: decode(tokens, position)
    Ctx->>Ctx: invalidateLogits()
    Note over Ctx: Detach old buffer<br/>_decodeStepId++
    Ctx->>CPP: llama_decode()
    CPP-->>Ctx: logits written to internal buffer
    Ctx-->>JS: success

    Note over JS,CPP: First getLogits() Call
    JS->>Ctx: getLogits()
    Ctx->>Ctx: Check memoization<br/>(_logitsStepId == _decodeStepId?)
    Note over Ctx: No existing buffer
    Ctx->>CPP: lloyal::logits::get(ctx, -1)
    CPP-->>Ctx: float* logits
    Ctx->>Buf: Create ArrayBuffer (zero-copy)
    Ctx->>Ctx: Store reference<br/>_logitsBufferRef = buffer<br/>_logitsStepId = _decodeStepId
    Ctx-->>JS: Float32Array view

    Note over JS,CPP: Second getLogits() Call (Same Step)
    JS->>Ctx: getLogits()
    Ctx->>Ctx: Check memoization<br/>(_logitsStepId == _decodeStepId?)
    Note over Ctx: Buffer exists, same step
    Ctx-->>JS: Return existing buffer

    Note over JS,CPP: Next Decode Invalidates Buffer
    JS->>Ctx: decode(nextTokens, position)
    Ctx->>Ctx: invalidateLogits()
    Ctx->>Buf: Detach ArrayBuffer
    Note over Buf: Buffer detached<br/>byteLength = 0
    Ctx->>Ctx: _decodeStepId++
    Ctx->>CPP: llama_decode()
    CPP-->>Ctx: new logits
    Ctx-->>JS: success

    Note over JS,Buf: Old Buffer Access (Error)
    JS->>Buf: Access old logits[0]
    Buf-->>JS: TypeError (detached)

Loading

[greptile-apps]

_{6 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

logic: comment says "weak reference" but reference count is 1 (strong reference). should use 0 for weak reference to avoid preventing garbage collection

Suggested change

	_logitsBufferRef = Napi::Reference<Napi::ArrayBuffer>::New(buffer, 1);
	_logitsBufferRef = Napi::Reference<Napi::ArrayBuffer>::New(buffer, 0);

Prompt To Fix With AI

This is a comment left during a code review.
Path: src/SessionContext.cpp
Line: 699:699

Comment:
**logic:** comment says "weak reference" but reference count is 1 (strong reference). should use `0` for weak reference to avoid preventing garbage collection

```suggestion
  _logitsBufferRef = Napi::Reference<Napi::ArrayBuffer>::New(buffer, 0);
```

How can I resolve this? If you propose a fix, please make it concise.

[lloyal-research]

The comment in the header says "Weak reference to detach on revocation" but that's describing the intent (we hold it so we can detach it later), not the N-API reference semantics.

Why we need refcount=1 (strong reference):

• We call buffer.Detach() later in invalidateLogits()
• If refcount=0, the buffer could be GC'd before we call Detach
• We need to keep it alive so we can explicitly invalidate it

Why refcount=0 would break things:

• Buffer gets GC'd
• _logitsBufferRef.Value() returns invalid/detached buffer
• buffer.Detach() fails or crashes

The fix is to update the misleading comment in the header, not change the refcount