Skip to content

Commit 8bbbc16

Browse files
vicaryvicary
committed
fix: address PR review comments - canonical ID resolution, config whitespace, recall text deduplication
1 parent d87a28e commit 8bbbc16

4 files changed

Lines changed: 80 additions & 31 deletions

File tree

‎src/config.test.ts‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,28 @@ describe("config", () => {
224224
);
225225
});
226226

227+
it("accepts endpoint-like config values with incidental surrounding whitespace", () => {
228+
setConfigExplorerAdapterForTesting(() =>
229+
makeAdapter({
230+
searchResult: {
231+
endpoint: " http://legacy.example/mcp ",
232+
redis: {
233+
endpoint: " redis://trimmed:6379 ",
234+
},
235+
graphiti: {
236+
endpoint: " http://nested.example/mcp ",
237+
},
238+
},
239+
})
240+
);
241+
242+
const config = loadConfig();
243+
244+
assertEquals(config.endpoint, "http://nested.example/mcp");
245+
assertEquals(config.graphiti.endpoint, "http://nested.example/mcp");
246+
assertEquals(config.redis.endpoint, "redis://trimmed:6379");
247+
});
248+
227249
it("fails open to defaults when config discovery search fails", () => {
228250
using _homedir = stub(os, "homedir", () => "/users/tester");
229251
setConfigExplorerAdapterForTesting(() =>

‎src/config.ts‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,14 @@ const readString = (
6767
): string | undefined =>
6868
typeof value[key] === "string" ? value[key] as string : undefined;
6969

70+
const readTrimmedString = (
71+
value: Record<string, unknown>,
72+
key: string,
73+
): string | undefined => {
74+
const entry = readString(value, key);
75+
return entry?.trim() || undefined;
76+
};
77+
7078
const readNumber = (
7179
value: Record<string, unknown>,
7280
key: string,
@@ -82,14 +90,14 @@ const normalizeConfig = (value: unknown): RawGraphitiConfig => {
8290
) as Partial<T>;
8391

8492
const config: RawGraphitiConfig = {
85-
endpoint: readString(value, "endpoint"),
93+
endpoint: readTrimmedString(value, "endpoint"),
8694
groupIdPrefix: readString(value, "groupIdPrefix"),
8795
driftThreshold: readNumber(value, "driftThreshold"),
8896
};
8997

9098
if (isRecord(value.redis)) {
9199
config.redis = compact({
92-
endpoint: readString(value.redis, "endpoint"),
100+
endpoint: readTrimmedString(value.redis, "endpoint"),
93101
batchSize: readNumber(value.redis, "batchSize"),
94102
batchMaxBytes: readNumber(value.redis, "batchMaxBytes"),
95103
sessionTtlSeconds: readNumber(value.redis, "sessionTtlSeconds"),
@@ -100,7 +108,7 @@ const normalizeConfig = (value: unknown): RawGraphitiConfig => {
100108

101109
if (isRecord(value.graphiti)) {
102110
config.graphiti = compact({
103-
endpoint: readString(value.graphiti, "endpoint"),
111+
endpoint: readTrimmedString(value.graphiti, "endpoint"),
104112
groupIdPrefix: readString(value.graphiti, "groupIdPrefix"),
105113
driftThreshold: readNumber(value.graphiti, "driftThreshold"),
106114
});
@@ -117,7 +125,7 @@ const isUnitInterval = (value: number | undefined): value is number =>
117125
value <= 1;
118126

119127
const isValidUrlString = (value: string | undefined): value is string => {
120-
if (!value?.trim()) return false;
128+
if (!value) return false;
121129
try {
122130
new URL(value);
123131
return true;

‎src/handlers/event.ts‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -330,9 +330,12 @@ export function createEventHandler(deps: EventHandlerDeps): EventHook {
330330
if (event.type === "message.part.updated") {
331331
const part = event.properties.part;
332332
if (!isTextPart(part)) return true;
333-
const canonicalSessionId = await sessionManager.resolveCanonicalSessionId(
334-
part.sessionID,
335-
) ?? part.sessionID;
333+
const {
334+
state,
335+
resolved,
336+
canonicalSessionId: resolvedCanonicalSessionId,
337+
} = await sessionManager.resolveSessionState(part.sessionID);
338+
const canonicalSessionId = resolvedCanonicalSessionId ?? part.sessionID;
336339
sessionManager.markResolvedSessionActive(
337340
part.sessionID,
338341
canonicalSessionId,
@@ -352,9 +355,6 @@ export function createEventHandler(deps: EventHandlerDeps): EventHook {
352355
return true;
353356
}
354357

355-
const { state, resolved } = await sessionManager.resolveSessionState(
356-
part.sessionID,
357-
);
358358
if (!resolved || !state?.isMain) return true;
359359

360360
const assistantText = sessionManager.finalizeAssistantMessage(

‎src/services/batch-drain.ts‎

Lines changed: 40 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -33,11 +33,24 @@ class DrainClaimLostError extends Error {
3333
const makeBatchKey = (entries: DrainQueueEntry[]): string =>
3434
`${entries[0]?.event.id ?? "empty"}:${entries.at(-1)?.event.id ?? "empty"}`;
3535

36-
const getDrainableEntryIds = (entries: DrainQueueEntry[]): Set<string> => {
36+
type PreparedDrainEntry = {
37+
entry: DrainQueueEntry;
38+
recallText: string;
39+
};
40+
41+
const prepareDrainEntries = (
42+
entries: DrainQueueEntry[],
43+
): PreparedDrainEntry[] =>
44+
entries.map((entry) => ({
45+
entry,
46+
recallText: getDrainEntryRecallText(entry),
47+
}));
48+
49+
const getDrainableEntryIds = (entries: PreparedDrainEntry[]): Set<string> => {
3750
const drainableEntryIds = new Set<string>();
3851
for (const entry of entries) {
3952
if (shouldDrainEntry(entry)) {
40-
drainableEntryIds.add(entry.event.id);
53+
drainableEntryIds.add(entry.entry.event.id);
4154
}
4255
}
4356
return drainableEntryIds;
@@ -46,40 +59,44 @@ const getDrainableEntryIds = (entries: DrainQueueEntry[]): Set<string> => {
4659
const getDrainEntryRecallText = (entry: DrainQueueEntry): string =>
4760
sanitizeMemoryInput(getSessionEventRecallText(entry.event));
4861

49-
const buildGraphitiEpisodeBody = (entry: DrainQueueEntry): string => {
50-
const refs = entry.event.refs?.length
51-
? `\nRefs: ${entry.event.refs.join(", ")}`
62+
const buildGraphitiEpisodeBody = (entry: PreparedDrainEntry): string => {
63+
const refs = entry.entry.event.refs?.length
64+
? `\nRefs: ${entry.entry.event.refs.join(", ")}`
5265
: "";
53-
const keywords = entry.event.keywords?.length
54-
? `\nKeywords: ${entry.event.keywords.join(", ")}`
66+
const keywords = entry.entry.event.keywords?.length
67+
? `\nKeywords: ${entry.entry.event.keywords.join(", ")}`
5568
: "";
5669
return sanitizeMemoryInput(
5770
[
58-
`Category: ${entry.event.category}`,
59-
`Role: ${entry.event.role}`,
60-
`Summary: ${entry.event.summary}`,
61-
entry.event.detail ? `Detail: ${entry.event.detail}` : "",
62-
entry.event.continuityText
63-
? `Continuity: ${entry.event.continuityText}`
64-
: getDrainEntryRecallText(entry),
71+
`Category: ${entry.entry.event.category}`,
72+
`Role: ${entry.entry.event.role}`,
73+
`Summary: ${entry.entry.event.summary}`,
74+
entry.entry.event.detail ? `Detail: ${entry.entry.event.detail}` : "",
75+
entry.entry.event.continuityText
76+
? `Continuity: ${entry.entry.event.continuityText}`
77+
: entry.recallText,
6578
keywords,
6679
refs,
6780
].filter(Boolean).join("\n"),
6881
);
6982
};
7083

71-
const shouldDrainEntry = (entry: DrainQueueEntry): boolean => {
72-
const text = getDrainEntryRecallText(entry);
84+
const shouldDrainEntry = (entry: PreparedDrainEntry): boolean => {
85+
const text = entry.recallText;
7386
if (!text) return false;
7487
if (looksLikeToolTranscript(text)) return false;
7588
if (looksLikeOperationalChatter(text)) return false;
7689
if (looksTranscriptHeavy(text)) return false;
7790
if (
78-
entry.event.role === "assistant" && entry.event.category !== "discovery"
91+
entry.entry.event.role === "assistant" &&
92+
entry.entry.event.category !== "discovery"
7993
) {
8094
return false;
8195
}
82-
if (entry.event.category === "message" && entry.event.role !== "user") {
96+
if (
97+
entry.entry.event.category === "message" &&
98+
entry.entry.event.role !== "user"
99+
) {
83100
return false;
84101
}
85102
return true;
@@ -175,9 +192,10 @@ export class BatchDrainService {
175192
}
176193

177194
const batch = claimed.entries;
195+
const preparedBatch = prepareDrainEntries(batch);
178196
const batchKey = makeBatchKey(batch);
179197
const eventIds = batch.map((entry) => entry.event.id);
180-
const drainableEntryIds = getDrainableEntryIds(batch);
198+
const drainableEntryIds = getDrainableEntryIds(preparedBatch);
181199
if (drainableEntryIds.size === 0) {
182200
await this.events.markBatchSuccess(groupId, claimed.claimToken, batch);
183201
await this.redis.deleteKey(drainRetryKey(groupId, batchKey));
@@ -241,12 +259,13 @@ export class BatchDrainService {
241259
let checkpointedCount = 0;
242260

243261
try {
244-
for (const entry of batch) {
262+
for (const preparedEntry of preparedBatch) {
263+
const entry = preparedEntry.entry;
245264
if (drainableEntryIds.has(entry.event.id)) {
246265
await assertClaimOwnership();
247266
await graphiti.addMemory({
248267
name: `${entry.event.category}:${entry.event.id}`,
249-
episodeBody: buildGraphitiEpisodeBody(entry),
268+
episodeBody: buildGraphitiEpisodeBody(preparedEntry),
250269
groupId,
251270
source: "text",
252271
sourceDescription: `session-event:${entry.event.category}`,

0 commit comments

Comments
 (0)