fix(tess): preserve redaction across stream chunks
This commit is contained in:
@@ -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);
|
||||||
});
|
});
|
||||||
|
|
||||||
|
|||||||
@@ -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;
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user