Skip to content

Commit 7d76bbf

Browse files
committed
fix(truncate): unique cache keys per payload + read_cache truncation observability
Follow-ups from PR #480 review: 1. Multi-block cache-key collision. cache.GenerateKey mixed the timestamp in at second granularity. forwardContentResult truncates each oversized TextContent block independently and the truncator derives the key from toolName+args+time.Now(); two blocks truncated in the same wall-clock second produced the SAME key, so the second Store overwrote the first and the first banner resolved to the wrong payload. Same hazard for rapid recursive read_cache truncation. Now nanosecond granularity. 2. Doc wording. maybeTruncateAndCacheText doc now states termination is agent-driven (not strictly algorithmic) and that successive identical calls are NOT idempotent (distinct timestamped keys), bounded only by cache TTL + cleanup. 3. read_cache truncation observability. handleReadCache discarded wasTruncated; it has no token-metrics plumbing (reports via activity log). A recursively re-truncated page was invisible to operators — now surfaced as a structured Info log with key/offset/limit/size. Tests: new GenerateKey sub-second regression assertion; TestForwardContentResult_MultipleTextBlocksDistinctKeys. Full internal/server, internal/cache, internal/truncate pass (incl. -race).
1 parent 95d09b0 commit 7d76bbf

5 files changed

Lines changed: 102 additions & 9 deletions

File tree

internal/cache/manager.go

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,11 +62,20 @@ func NewManager(db *bbolt.DB, logger *zap.Logger) (*Manager, error) {
6262
return manager, nil
6363
}
6464

65-
// GenerateKey generates a cache key from tool name, arguments, and timestamp
65+
// GenerateKey generates a cache key from tool name, arguments, and timestamp.
66+
//
67+
// The timestamp is mixed in at nanosecond granularity. The truncator calls
68+
// this with time.Now() once per truncated payload; a single upstream result
69+
// can carry multiple oversized TextContent blocks, and recursive read_cache
70+
// truncation can mint several keys in quick succession. At second granularity
71+
// any of those that landed in the same wall-clock second collided on one key,
72+
// so a later Store silently overwrote an earlier payload and the earlier
73+
// banner resolved to the wrong data. Nanosecond granularity makes each call's
74+
// key unique in practice.
6675
func GenerateKey(toolName string, args map[string]interface{}, timestamp time.Time) string {
6776
// Create a consistent string representation
6877
argsJSON, _ := json.Marshal(args)
69-
input := fmt.Sprintf("%s:%s:%d", toolName, string(argsJSON), timestamp.Unix())
78+
input := fmt.Sprintf("%s:%s:%d", toolName, string(argsJSON), timestamp.UnixNano())
7079

7180
hash := sha256.Sum256([]byte(input))
7281
return hex.EncodeToString(hash[:])

internal/cache/manager_test.go

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,17 @@ func TestGenerateKey(t *testing.T) {
8080
if len(key1) != 64 { // SHA256 hex = 64 chars
8181
t.Errorf("Expected key length 64, got %d", len(key1))
8282
}
83+
84+
// Sub-second-distinct timestamps must produce different keys. The
85+
// truncator calls GenerateKey with time.Now() for every truncated block;
86+
// at second granularity two blocks truncated in the same wall-clock
87+
// second collided on one key, so the second Store overwrote the first and
88+
// read_cache for the first banner resolved to the wrong payload.
89+
tsA := time.Unix(1700000000, 0)
90+
tsB := tsA.Add(time.Nanosecond)
91+
if GenerateKey("test_tool", args1, tsA) == GenerateKey("test_tool", args1, tsB) {
92+
t.Error("Timestamps 1ns apart must produce different keys (multi-block collision regression)")
93+
}
8394
}
8495

8596
func TestStoreAndGet(t *testing.T) {

internal/server/content_forward.go

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -144,12 +144,21 @@ func forwardContentResult(result interface{}, truncator *truncate.Truncator, cac
144144
// the returned snippet carries the standard "use read_cache" banner that
145145
// points at it.
146146
//
147-
// The paginableUnits parameter is a guard against infinite recursion for
148-
// built-in tools that may themselves be paginating (notably read_cache). If
149-
// paginableUnits <= 1 the caller has nothing further to subdivide — caching
150-
// a fresh key under those circumstances just produces an identical record
151-
// the agent can never escape from. In that case the text is returned as-is
152-
// (uncaught, exceeding the limit) and the caller decides what to do.
147+
// The paginableUnits parameter is the recursion guard for built-in tools that
148+
// may themselves be paginating (notably read_cache). If paginableUnits <= 1
149+
// the caller has nothing further to subdivide — caching a fresh key here would
150+
// just hand the agent a new key resolving to the same oversize payload, an
151+
// inescapable loop. In that case the text is returned as-is (still over the
152+
// limit) and the caller decides what to do.
153+
//
154+
// Termination is therefore agent-driven, not strictly algorithmic: a level
155+
// whose records all fit individually but overflow collectively keeps minting
156+
// a fresh key per call until the agent pages down to the single-record
157+
// short-circuit. Each call's key is unique (timestamped at nanosecond
158+
// granularity), so successive identical calls are NOT idempotent — they
159+
// produce distinct cache entries. Unbounded growth is bounded only by the
160+
// cache's TTL + background cleanup (see cache.DefaultTTL), which is acceptable
161+
// because truncation is the exceptional path.
153162
//
154163
// Returns:
155164
// - the text to actually emit (possibly truncated)

internal/server/content_forward_test.go

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -287,6 +287,55 @@ func TestForwardContentResult_RoundTripViaRealCacheManager(t *testing.T) {
287287
assert.Equal(t, float64(19), idOf(page2.Records[9]), "page 2 should end at id 19")
288288
}
289289

290+
// TestForwardContentResult_MultipleTextBlocksDistinctKeys is a regression test
291+
// for the multi-block cache-key collision. When an upstream result carries
292+
// more than one oversized TextContent block, each block is truncated
293+
// independently and gets its own banner + cache key. The truncator derives the
294+
// key from toolName+args+timestamp; at the previous second-granularity, two
295+
// blocks truncated within the same wall-clock second produced the SAME key, so
296+
// the second Store overwrote the first and the first banner resolved to the
297+
// wrong block's payload. Each block must persist under a distinct key and that
298+
// key must resolve to that block's own content.
299+
func TestForwardContentResult_MultipleTextBlocksDistinctKeys(t *testing.T) {
300+
mkBig := func(tag string) string {
301+
recs := make([]string, 0, 40)
302+
for i := 0; i < 40; i++ {
303+
recs = append(recs, fmt.Sprintf(`{"id":%d,"tag":"%s","pad":"%s"}`, i, tag, strings.Repeat(tag, 20)))
304+
}
305+
return `[` + strings.Join(recs, ",") + `]`
306+
}
307+
blockA := mkBig("A")
308+
blockB := mkBig("B")
309+
310+
upstream := &mcp.CallToolResult{
311+
Content: []mcp.Content{
312+
mcp.NewTextContent(blockA),
313+
mcp.NewTextContent(blockB),
314+
},
315+
}
316+
store := &captureStore{}
317+
_, _, wasTruncated := forwardContentResult(
318+
upstream,
319+
truncate.NewTruncator(500),
320+
store,
321+
nil,
322+
"github:pull_request_read",
323+
map[string]interface{}{"perPage": 100},
324+
)
325+
require.True(t, wasTruncated)
326+
327+
store.mu.Lock()
328+
defer store.mu.Unlock()
329+
require.Len(t, store.calls, 2, "each oversized text block must produce its own Store call")
330+
331+
k0, k1 := store.calls[0].key, store.calls[1].key
332+
assert.NotEmpty(t, k0)
333+
assert.NotEmpty(t, k1)
334+
assert.NotEqual(t, k0, k1, "distinct text blocks must persist under distinct cache keys")
335+
assert.Equal(t, blockA, store.calls[0].content, "block A's key must resolve to block A's content")
336+
assert.Equal(t, blockB, store.calls[1].content, "block B's key must resolve to block B's content")
337+
}
338+
290339
// TestMaybeTruncateAndCacheText_HappyPathStores asserts that an oversized text
291340
// with more than one paginable unit gets truncated AND stored under the
292341
// embedded cache key — the contract required for read_cache pagination to keep

internal/server/mcp.go

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4129,7 +4129,7 @@ func (p *MCPProxyServer) handleReadCache(ctx context.Context, request mcp.CallTo
41294129
// in that case there's nothing this layer can subdivide, so the oversize
41304130
// text flows through unchanged. p.logger receives a zap.Warn if the cache
41314131
// write fails so the resulting "cache key not found" is diagnosable.
4132-
text, _ := maybeTruncateAndCacheText(
4132+
text, reTruncated := maybeTruncateAndCacheText(
41334133
string(jsonResult),
41344134
"read_cache",
41354135
args,
@@ -4139,6 +4139,21 @@ func (p *MCPProxyServer) handleReadCache(ctx context.Context, request mcp.CallTo
41394139
p.logger,
41404140
)
41414141

4142+
// read_cache has no token-metrics plumbing (it reports via the activity
4143+
// log, not the tokenMetrics/toolCallRecord path used by upstream tool
4144+
// calls). A recursively re-truncated page is otherwise invisible to
4145+
// operators, so surface it as a structured log: it signals the agent is
4146+
// walking a payload deep enough that even one cache page overflows.
4147+
if reTruncated {
4148+
p.logger.Info("read_cache output exceeded the response limit and was recursively truncated-and-cached",
4149+
zap.String("requested_key", key),
4150+
zap.Int("offset", offset),
4151+
zap.Int("limit", limit),
4152+
zap.Int("page_records", len(response.Records)),
4153+
zap.Int("full_bytes", len(jsonResult)),
4154+
)
4155+
}
4156+
41424157
// Spec 024: Emit success event with args and response
41434158
p.emitActivityInternalToolCall("read_cache", "", "", "", sessionID, requestID, "success", "", time.Since(startTime).Milliseconds(), args, response, nil, "")
41444159

0 commit comments

Comments
 (0)