fix(tess): preserve redaction across stream chunks

This commit is contained in:
Jarvis
2026-07-12 19:32:58 -05:00
parent 98b4f95284
commit 70851d9474
2 changed files with 122 additions and 14 deletions

View File

@@ -2,7 +2,7 @@ import { describe, expect, it, vi } from 'vitest';
import { ChatGateway } from './chat.gateway.js'; import { ChatGateway } from './chat.gateway.js';
const CONVERSATION_ID = 'conversation-1'; const CONVERSATION_ID = 'conversation-1';
const CANARY = 'token=sk_canary12345678'; const CANARY = 'sk_canary12345678';
type GatewayInternals = { type GatewayInternals = {
clientSessions: Map<string, unknown>; clientSessions: Map<string, unknown>;
@@ -31,7 +31,7 @@ function buildGateway() {
} }
describe('ChatGateway redaction boundary', (): void => { describe('ChatGateway redaction boundary', (): void => {
it('redacts an assistant delta before egress and before its persistence buffer', (): void => { it('redacts a secret split across assistant deltas before egress and persistence', (): void => {
const { gateway } = buildGateway(); const { gateway } = buildGateway();
const client = { const client = {
connected: true, connected: true,
@@ -51,14 +51,21 @@ describe('ChatGateway redaction boundary', (): void => {
gateway.relayEvent(client, CONVERSATION_ID, { gateway.relayEvent(client, CONVERSATION_ID, {
type: 'message_update', type: 'message_update',
assistantMessageEvent: { type: 'text_delta', delta: CANARY }, assistantMessageEvent: { type: 'text_delta', delta: 'sk_canary' },
});
expect(JSON.stringify(client.emit.mock.calls)).not.toContain('sk_canary');
gateway.relayEvent(client, CONVERSATION_ID, {
type: 'message_update',
assistantMessageEvent: { type: 'text_delta', delta: '12345678 ' },
}); });
expect(client.emit).toHaveBeenCalledWith('agent:text', { expect(client.emit).toHaveBeenCalledWith('agent:text', {
conversationId: CONVERSATION_ID, conversationId: CONVERSATION_ID,
text: '[REDACTED_SECRET] ', text: '[REDACTED_SECRET] ',
}); });
expect(session.assistantText).toBe('[REDACTED_SECRET]'); expect(session.assistantText).toBe(`${CANARY} `);
expect(JSON.stringify(client.emit.mock.calls)).not.toContain(CANARY); expect(JSON.stringify(client.emit.mock.calls)).not.toContain(CANARY);
}); });

View File

@@ -108,6 +108,9 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
private readonly logger = new Logger(ChatGateway.name); private readonly logger = new Logger(ChatGateway.name);
private readonly clientSessions = new Map<string, ClientSession>(); private readonly clientSessions = new Map<string, ClientSession>();
/** Raw stream fragments are kept in memory only until they are safe to redact and emit. */
private readonly textEgressBuffers = new Map<string, string>();
private readonly thinkingEgressBuffers = new Map<string, string>();
private readonly discordReplayProtector = new DiscordReplayProtector(); private readonly discordReplayProtector = new DiscordReplayProtector();
constructor( constructor(
@@ -156,6 +159,8 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
); );
this.clientSessions.delete(client.id); this.clientSessions.delete(client.id);
} }
this.textEgressBuffers.delete(client.id);
this.thinkingEgressBuffers.delete(client.id);
} }
private getClientScope(client: Socket): ActorTenantScope | null { private getClientScope(client: Socket): ActorTenantScope | null {
@@ -746,6 +751,81 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
} }
} }
private appendAndFlushRedactedEgress(
client: Socket,
conversationId: string,
eventName: 'agent:text' | 'agent:thinking',
buffers: Map<string, string>,
delta: string,
): void {
buffers.set(client.id, `${buffers.get(client.id) ?? ''}${delta}`);
this.flushRedactedEgress(client, conversationId, eventName, buffers, false);
}
/**
* Holds any suffix that could become a secret, email, or phone number after a
* later stream chunk. This avoids relying on downstream redaction after data
* has already reached the socket.
*/
private flushRedactedEgress(
client: Socket,
conversationId: string,
eventName: 'agent:text' | 'agent:thinking',
buffers: Map<string, string>,
final: boolean,
): void {
const buffered = buffers.get(client.id) ?? '';
const releaseLength = final ? buffered.length : this.safeRedactionPrefixLength(buffered);
const released = buffered.slice(0, releaseLength);
const pending = buffered.slice(releaseLength);
if (pending) {
buffers.set(client.id, pending);
} else {
buffers.delete(client.id);
}
if (released) {
client.emit(eventName, {
conversationId,
text: redactSensitiveContent(released).content,
});
}
}
private safeRedactionPrefixLength(content: string): number {
let retainedFrom = content.length;
// Retain the current token because it may become a split secret or email.
const token = /(?:^|\s)(\S*)$/.exec(content);
if (token) {
const matched = token[0] ?? '';
const trailingToken = token[1] ?? '';
retainedFrom = token.index + matched.length - trailingToken.length;
}
// The secret classifier accepts whitespace around ':' and '=', so preserve
// a pending label until its value and delimiter are both complete.
const secretLabel = /(?:api[_-]?key|token|password|secret)\s*[:=]\s*$/i.exec(content);
if (secretLabel) {
retainedFrom = Math.min(retainedFrom, secretLabel.index);
}
// Phone numbers can contain whitespace and punctuation; preserve the full
// trailing numeric candidate until a non-phone character establishes a boundary.
const phone = /(?:^|[^A-Za-z0-9_])(\+?\d[\d(). -]*)$/.exec(content);
if (phone) {
const matched = phone[0] ?? '';
const trailingPhoneCandidate = phone[1] ?? '';
retainedFrom = Math.min(
retainedFrom,
phone.index + matched.length - trailingPhoneCandidate.length,
);
}
return retainedFrom;
}
private relayEvent(client: Socket, conversationId: string, event: AgentSessionEvent): void { private relayEvent(client: Socket, conversationId: string, event: AgentSessionEvent): void {
if (!client.connected) { if (!client.connected) {
this.logger.warn( this.logger.warn(
@@ -763,6 +843,8 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
cs.toolCalls = []; cs.toolCalls = [];
cs.pendingToolCalls.clear(); cs.pendingToolCalls.clear();
} }
this.textEgressBuffers.set(client.id, '');
this.thinkingEgressBuffers.set(client.id, '');
client.emit('agent:start', { conversationId }); client.emit('agent:start', { conversationId });
break; break;
} }
@@ -791,6 +873,20 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
} }
: undefined; : undefined;
this.flushRedactedEgress(
client,
conversationId,
'agent:text',
this.textEgressBuffers,
true,
);
this.flushRedactedEgress(
client,
conversationId,
'agent:thinking',
this.thinkingEgressBuffers,
true,
);
client.emit('agent:end', { client.emit('agent:end', {
conversationId, conversationId,
usage: usagePayload, usage: usagePayload,
@@ -859,21 +955,26 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
case 'message_update': { case 'message_update': {
const assistantEvent = event.assistantMessageEvent; const assistantEvent = event.assistantMessageEvent;
if (assistantEvent.type === 'text_delta') { if (assistantEvent.type === 'text_delta') {
// Accumulate assistant text for persistence // Keep raw stream material in memory only; persist and emit only redacted text.
const cs = this.clientSessions.get(client.id); const cs = this.clientSessions.get(client.id);
const redactedDelta = redactSensitiveContent(assistantEvent.delta).content;
if (cs) { if (cs) {
cs.assistantText += redactedDelta; cs.assistantText += assistantEvent.delta;
} }
client.emit('agent:text', { this.appendAndFlushRedactedEgress(
client,
conversationId, conversationId,
text: redactedDelta, 'agent:text',
}); this.textEgressBuffers,
assistantEvent.delta,
);
} else if (assistantEvent.type === 'thinking_delta') { } else if (assistantEvent.type === 'thinking_delta') {
client.emit('agent:thinking', { this.appendAndFlushRedactedEgress(
client,
conversationId, conversationId,
text: redactSensitiveContent(assistantEvent.delta).content, 'agent:thinking',
}); this.thinkingEgressBuffers,
assistantEvent.delta,
);
} }
break; break;
} }