From d6515c30d8fcd310a19ea667337ed97978e4622f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 27 Oct 2026 23:51:31 +0000 Subject: [PATCH] Forward E2EE tool calls to OpenAI clients The proxy kept only decrypted text and always finished with stop, then rejected the follow-up because an assistant tool call uses content null. Stream tool_calls deltas in order, assemble them when the client does not stream, and send tool-result bodies to encrypt while leaving call metadata in the clear. The library pin is the unpublished commit 7b4e652 until 0.3.0 is on npm. Co-authored-by: guil --- ARCHITECTURE.md | 13 +- README.md | 2 +- package-lock.json | 14 +- package.json | 2 +- src/chat/chat.service.ts | 60 ++++- src/chat/dto/chat-completion-response.dto.ts | 17 +- src/chat/dto/chat-message.dto.ts | 34 ++- src/common/interfaces/openai.interface.ts | 34 ++- src/common/utils/openai.util.ts | 127 +++++++++-- src/e2ee/e2ee.service.ts | 21 +- test/helpers/test-app.ts | 24 +- test/openai-compliance.spec.ts | 223 ++++++++++++++++++- test/openai.util.spec.ts | 87 +++++++- test/tool-encryption.spec.ts | 87 ++++++++ 14 files changed, 667 insertions(+), 78 deletions(-) create mode 100644 test/tool-encryption.spec.ts diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 8da3be5..fa9e802 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -90,12 +90,12 @@ The proxy forwards the request to Venice unchanged. Non-streaming returns the JS ### E2EE, streaming (`stream: true`) 1. `E2eeService.createSession(model)` — fetches the TEE attestation, **verifies it**, and performs ECDH key exchange (cached for `SESSION_TTL_MS`). -2. `E2eeService.encrypt(messages, session)` — AES-256-GCM encrypts each message; returns `encryptedMessages`, TEE headers, and `venice_parameters: { enable_e2ee: true }`. -3. The proxy POSTs to Venice with `stream: true`. -4. `E2eeService.decryptStream(body, session)` decrypts each SSE chunk; the controller re-wraps the plaintext deltas as standard OpenAI `chat.completion.chunk` events, terminated by `data: [DONE]`. +2. `E2eeService.encrypt(messages, session)` — AES-256-GCM encrypts prompt and tool-result `content`. `tool_calls`, `tool_call_id`, and `name` stay JSON. `content: null` on a tool-only assistant turn stays `null`. Returns `encryptedMessages`, TEE headers, and `venice_parameters: { enable_e2ee: true }`. +3. The proxy POSTs to Venice with `stream: true`. `tools` and `tool_choice` are forwarded as-is. +4. `E2eeService.decryptStream(body, session)` yields decrypted text, unchanged `tool_calls` deltas, and the real `finish_reason`. The controller re-wraps those as OpenAI `chat.completion.chunk` events, terminated by `data: [DONE]`. ### E2EE, non-streaming (`stream: false` or omitted) -Venice **requires** streaming for E2EE. So the proxy still streams upstream and decrypts, but **buffers** the deltas and returns a single OpenAI `chat.completion` object. This is why non-streaming E2EE requests still work for clients that don't speak SSE — the streaming requirement is an upstream constraint we hide. +Venice **requires** streaming for E2EE. So the proxy still streams upstream and decrypts, but **buffers** text and tool-call deltas (assembled by `index`) and returns a single OpenAI `chat.completion` object. A tool-only turn has `content: null` and `finish_reason: "tool_calls"`. This is why non-streaming E2EE requests still work for clients that don't speak SSE — the streaming requirement is an upstream constraint we hide. --- @@ -129,10 +129,11 @@ Proxy (client side) Venice CPU TEE (Intel TDX) Two subtleties worth recording: -- **Per-chunk ephemeral keys.** Every response chunk carries a *fresh* server ephemeral public key, so each chunk requires its own ECDH derivation. The library handles this in `decryptStream`; we just consume plaintext deltas. +- **Per-chunk ephemeral keys.** Every encrypted response chunk carries a *fresh* server ephemeral public key, so each chunk requires its own ECDH derivation. The library handles this in `decryptStream`; we consume decrypted text plus plaintext tool-call deltas. - **Plaintext passthrough.** Some tokens (e.g. whitespace) may arrive unencrypted; the library handles that and the `[DONE]` sentinel internally. +- **Tool metadata is plaintext on the wire.** Tool names, JSON schemas, call ids, and arguments are not part of the encrypted `content` channel — the enclave chat template reads them as JSON. Prompts and tool-result bodies are ciphertext. -Plaintext exists in only two places: **on this proxy** (briefly, in memory) and **inside the TEE**. It never traverses the network in the clear. +Prompt and tool-result plaintext exists in only two places: **on this proxy** (briefly, in memory) and **inside the TEE**. That text never traverses the network in the clear. Tool-call metadata does. --- diff --git a/README.md b/README.md index c9b52aa..35b6670 100644 --- a/README.md +++ b/README.md @@ -166,7 +166,7 @@ Interactive API documentation is served at **[http://localhost:4337/docs](http:/ └─────────┘ plaintext SSE └──────────────────┘ encrypted stream └──────────────┘ ``` -- **`e2ee-*` models** → the proxy creates a verified TEE session, encrypts your messages, forwards them to Venice with `venice_parameters: { enable_e2ee: true }`, and decrypts the streamed response on the fly. +- **`e2ee-*` models** → the proxy creates a verified TEE session, encrypts prompts and tool-result bodies, forwards them to Venice with `venice_parameters: { enable_e2ee: true }`, and decrypts the streamed response on the fly. Tool names, schemas, call ids, and arguments stay plaintext (the enclave reads them as JSON). A tool call is returned as `delta.tool_calls` with `finish_reason: "tool_calls"`, so an agent can run the tool and send the result back without the history being flattened to text. - **All other models** → the request is proxied straight through to Venice unchanged. **Under the hood:** NestJS (TypeScript, strict mode); encryption via the [`@axlabs/venice-e2ee`](https://github.com/AxLabs/venice-e2ee) library (ECDH secp256k1 → HKDF-SHA256 → AES-256-GCM + TEE attestation); Swagger UI at `/docs`. diff --git a/package-lock.json b/package-lock.json index 53adda2..cdd5ca7 100644 --- a/package-lock.json +++ b/package-lock.json @@ -9,7 +9,7 @@ "version": "1.0.1", "license": "Apache-2.0", "dependencies": { - "@axlabs/venice-e2ee": "^0.2.0", + "@axlabs/venice-e2ee": "git+https://github.com/AxLabs/venice-e2ee.git#7b4e6528f7f870402224bc887de06f0952f7f575", "@nestjs/common": "^11.0.0", "@nestjs/config": "^4.0.0", "@nestjs/core": "^11.0.0", @@ -22,7 +22,7 @@ "rxjs": "^7.8.1" }, "bin": { - "venice-e2ee-proxy": "dist/main.js" + "venice-e2ee-proxy": "bin/venice-e2ee-proxy.js" }, "devDependencies": { "@eslint/js": "^9.17.0", @@ -172,8 +172,8 @@ }, "node_modules/@axlabs/venice-e2ee": { "version": "0.2.0", - "resolved": "https://registry.npmjs.org/@axlabs/venice-e2ee/-/venice-e2ee-0.2.0.tgz", - "integrity": "sha512-s5Q+oryOBbRiWczAlPX7CAT2hI3B9asrV1YTAcCkZaEzDryog+HtI95ZnJ9g1a1PjqGcJOjh8a5SEWjLtcqByg==", + "resolved": "git+https://github.com/AxLabs/venice-e2ee.git#7b4e6528f7f870402224bc887de06f0952f7f575", + "integrity": "sha512-vSZW0eyZ0d0DNsYWQoIb8yrIZG+/6RrlJ/z5ms1NXxvr6HLAGC2XgqepwkKUUQ1Rb0yNMrIMdaj7Tu3LGWA5+w==", "license": "GPL-3.0-only", "dependencies": { "@noble/hashes": "^2.0.1", @@ -192,9 +192,9 @@ } }, "node_modules/@axlabs/venice-e2ee/node_modules/@noble/hashes": { - "version": "2.2.0", - "resolved": "https://registry.npmjs.org/@noble/hashes/-/hashes-2.2.0.tgz", - "integrity": "sha512-IYqDGiTXab6FniAgnSdZwgWbomxpy9FtYvLKs7wCUs2a8RkITG+DFGO1DM9cr+E3/RgADRpFjrKVaJ1z6sjtEg==", + "version": "2.4.0", + "resolved": "https://registry.npmjs.org/@noble/hashes/-/hashes-2.4.0.tgz", + "integrity": "sha512-X5XaVWZIBCT7HHZGm5I7ZQXDwLG+bGXuSrMQAW+7Zvl87h1kmc1ZB1VSRJcpUfoUrGQp4Fkoxm5kZ+Ms+aW+eA==", "license": "MIT", "engines": { "node": ">= 20.19.0" diff --git a/package.json b/package.json index ef3935f..a081dd2 100644 --- a/package.json +++ b/package.json @@ -59,7 +59,7 @@ "prepublishOnly": "npm run build" }, "dependencies": { - "@axlabs/venice-e2ee": "^0.2.0", + "@axlabs/venice-e2ee": "git+https://github.com/AxLabs/venice-e2ee.git#7b4e6528f7f870402224bc887de06f0952f7f575", "@nestjs/common": "^11.0.0", "@nestjs/config": "^4.0.0", "@nestjs/core": "^11.0.0", diff --git a/src/chat/chat.service.ts b/src/chat/chat.service.ts index c5b4aef..d3540e1 100644 --- a/src/chat/chat.service.ts +++ b/src/chat/chat.service.ts @@ -1,12 +1,19 @@ import { Injectable, InternalServerErrorException, Logger } from '@nestjs/common'; import { E2eeService } from '../e2ee/e2ee.service'; import { VeniceHttpService } from '../common/venice/venice-http.service'; -import type { ChatCompletion, ChatCompletionChunk } from '../common/interfaces/openai.interface'; +import type { + AssembledToolCall, + ChatCompletion, + ChatCompletionChunk, +} from '../common/interfaces/openai.interface'; import { + appendToolCallDeltas, buildChatCompletion, buildContentChunk, buildFinalChunk, buildRoleChunk, + buildToolCallsChunk, + compactToolCalls, countIncomingMessageChars, countOutgoingE2eeMessageChars, generateCompletionId, @@ -37,8 +44,8 @@ export class ChatService { } /** - * E2EE streaming: synthesize OpenAI SSE chunks from decrypted text deltas. - * E2EE models always stream upstream (a Venice requirement). + * E2EE streaming: synthesize OpenAI SSE chunks from decrypted text deltas + * and plaintext tool-call deltas. E2EE models always stream upstream. */ async *streamE2EE(dto: ChatCompletionRequestDto): AsyncGenerator { this.logger.debug(`E2EE stream: establishing session for "${dto.model}"`); @@ -94,43 +101,72 @@ export class ChatService { let decryptedChunks = 0; let decryptedChars = 0; + let toolCallDeltas = 0; + let finishReason = 'stop'; yield buildRoleChunk({ id, model: dto.model, created }); - for await (const text of this.e2ee.decryptStream(response.body, session)) { - if (text.length > 0) { - decryptedChunks += 1; - decryptedChars += text.length; - yield buildContentChunk({ id, model: dto.model, content: text, created }); + for await (const event of this.e2ee.decryptStream(response.body, session)) { + switch (event.kind) { + case 'content': + if (event.content.length > 0) { + decryptedChunks += 1; + decryptedChars += event.content.length; + yield buildContentChunk({ id, model: dto.model, content: event.content, created }); + } + break; + case 'tool_calls': + toolCallDeltas += event.tool_calls.length; + yield buildToolCallsChunk({ + id, + model: dto.model, + created, + toolCalls: event.tool_calls, + }); + break; + case 'finish': + finishReason = event.finish_reason; + break; + default: { + const unexpected: never = event; + throw new Error(`Unhandled E2EE stream event: ${JSON.stringify(unexpected)}`); + } } } - yield buildFinalChunk({ id, model: dto.model, created }); + yield buildFinalChunk({ id, model: dto.model, created, finishReason }); this.logger.debug( - `E2EE stream: decrypted ${decryptedChunks} chunk(s), ${decryptedChars} char(s) (id=${id})`, + `E2EE stream: decrypted ${decryptedChunks} chunk(s), ${decryptedChars} char(s), ` + + `${toolCallDeltas} tool-call delta(s), finish_reason=${finishReason} (id=${id})`, ); } /** * E2EE non-streaming: Venice requires streaming for E2EE, so we stream + decrypt - * upstream, buffer the deltas, and return a single OpenAI completion object. + * upstream, buffer text and tool-call deltas, and return one completion object. */ async completeE2EE(dto: ChatCompletionRequestDto): Promise { this.logger.debug(`E2EE non-streaming: buffering decrypted stream for "${dto.model}"`); let content = ''; let id = ''; let finishReason = 'stop'; + const toolCalls: AssembledToolCall[] = []; for await (const chunk of this.streamE2EE(dto)) { id = chunk.id; const choice = chunk.choices[0]; if (choice?.delta?.content) content += choice.delta.content; + if (choice?.delta?.tool_calls) appendToolCallDeltas(toolCalls, choice.delta.tool_calls); if (choice?.finish_reason) finishReason = choice.finish_reason; } - this.logger.debug(`E2EE non-streaming: assembled ${content.length} char(s) (id=${id})`); + const assembled = compactToolCalls(toolCalls); + this.logger.debug( + `E2EE non-streaming: assembled ${content.length} char(s), ${assembled.length} tool call(s) (id=${id})`, + ); return buildChatCompletion({ id: id || generateCompletionId(), model: dto.model, content, finishReason, + toolCalls: assembled, }); } diff --git a/src/chat/dto/chat-completion-response.dto.ts b/src/chat/dto/chat-completion-response.dto.ts index 3e4b01a..ead7521 100644 --- a/src/chat/dto/chat-completion-response.dto.ts +++ b/src/chat/dto/chat-completion-response.dto.ts @@ -1,11 +1,22 @@ -import { ApiProperty } from '@nestjs/swagger'; +import { ApiProperty, ApiPropertyOptional } from '@nestjs/swagger'; class ChatCompletionMessageDto { @ApiProperty({ example: 'assistant' }) role!: string; - @ApiProperty({ example: 'Hello! How can I help you today?' }) - content!: string; + @ApiProperty({ + example: 'Hello! How can I help you today?', + nullable: true, + description: 'Null when the assistant message is only a tool call.', + }) + content!: string | null; + + @ApiPropertyOptional({ + description: 'Tool calls assembled from the stream, when the model asked to call tools.', + type: 'array', + items: { type: 'object' }, + }) + tool_calls?: unknown[]; } class ChatCompletionChoiceDto { diff --git a/src/chat/dto/chat-message.dto.ts b/src/chat/dto/chat-message.dto.ts index ad5ce4f..1fb28f8 100644 --- a/src/chat/dto/chat-message.dto.ts +++ b/src/chat/dto/chat-message.dto.ts @@ -1,11 +1,12 @@ import { ApiProperty, ApiPropertyOptional } from '@nestjs/swagger'; -import { IsDefined, IsOptional, IsString } from 'class-validator'; -import type { ChatMessageContent } from '../../common/interfaces/openai.interface'; +import { IsArray, IsOptional, IsString } from 'class-validator'; +import type { ChatMessageContent, ToolCallDelta } from '../../common/interfaces/openai.interface'; /** * A single chat message. Mirrors the OpenAI message object. `content` accepts - * either a plain string or the structured array form; both are normalized to - * text before encryption. + * a string, the structured array form, or null when the assistant turn is only + * a tool call. Tool-result bodies are encrypted; tool names, ids, and + * arguments are forwarded in the clear. */ export class ChatMessageDto { @ApiProperty({ @@ -17,15 +18,34 @@ export class ChatMessageDto { role!: string; @ApiProperty({ - description: 'The message content. A string, or an array of content parts.', + description: + 'The message content. A string, an array of content parts, or null when an assistant message is only a tool call.', + nullable: true, oneOf: [{ type: 'string' }, { type: 'array', items: { type: 'object' } }], example: 'Hello!', }) - @IsDefined() - content!: ChatMessageContent; + @IsOptional() + content?: ChatMessageContent; @ApiPropertyOptional({ description: 'Optional name of the message author.' }) @IsOptional() @IsString() name?: string; + + @ApiPropertyOptional({ + description: 'Id of the tool call a `tool` message answers.', + }) + @IsOptional() + @IsString() + tool_call_id?: string; + + @ApiPropertyOptional({ + description: + 'Tool calls emitted by an assistant message. Names, ids, and arguments stay plaintext.', + type: 'array', + items: { type: 'object' }, + }) + @IsOptional() + @IsArray() + tool_calls?: ToolCallDelta[]; } diff --git a/src/common/interfaces/openai.interface.ts b/src/common/interfaces/openai.interface.ts index a288357..2e37b24 100644 --- a/src/common/interfaces/openai.interface.ts +++ b/src/common/interfaces/openai.interface.ts @@ -13,13 +13,39 @@ export interface ChatContentPart { [key: string]: unknown; } -/** OpenAI message content can be a plain string or an array of typed parts. */ +/** OpenAI message content can be a plain string, an array of typed parts, or null. */ export type ChatMessageContent = string | ChatContentPart[] | null; +/** + * A streaming tool-call delta, or the call object stored on an assistant message. + * Names, ids, and arguments stay plaintext; Venice's enclave reads them as JSON. + */ +export interface ToolCallDelta { + index?: number; + id?: string; + type?: string; + function?: { + name?: string; + arguments?: string; + }; +} + +/** Tool call assembled from streaming deltas for a non-streaming completion. */ +export interface AssembledToolCall { + id: string; + type: string; + function: { + name: string; + arguments: string; + }; +} + export interface ChatMessage { role: ChatRole; - content: ChatMessageContent; + content?: ChatMessageContent; name?: string; + tool_call_id?: string; + tool_calls?: ToolCallDelta[]; [key: string]: unknown; } @@ -27,7 +53,8 @@ export interface ChatCompletionChoice { index: number; message: { role: 'assistant'; - content: string; + content: string | null; + tool_calls?: AssembledToolCall[]; }; finish_reason: string | null; logprobs: null; @@ -55,6 +82,7 @@ export interface ChatCompletionChunkChoice { delta: { role?: 'assistant'; content?: string; + tool_calls?: ToolCallDelta[]; }; finish_reason: string | null; logprobs: null; diff --git a/src/common/utils/openai.util.ts b/src/common/utils/openai.util.ts index 613c248..7b47bc6 100644 --- a/src/common/utils/openai.util.ts +++ b/src/common/utils/openai.util.ts @@ -1,9 +1,11 @@ import { randomBytes } from 'node:crypto'; import type { + AssembledToolCall, ChatCompletion, ChatCompletionChunk, ChatMessage, ChatMessageContent, + ToolCallDelta, } from '../interfaces/openai.interface'; const E2EE_PREFIX = 'e2ee-'; @@ -25,23 +27,53 @@ export function unixSeconds(): number { /** * Flattens OpenAI message content (string or array-of-parts) into a single - * string. `@axlabs/venice-e2ee` encrypts `{ role, content: string }` pairs, so any - * structured content is collapsed to its text parts. + * string. `null` and missing content stay empty here; callers that encrypt + * decide whether that becomes `null` or ciphertext. */ -export function flattenContent(content: ChatMessageContent): string { +export function flattenContent(content: ChatMessageContent | undefined): string { if (content === null || content === undefined) return ''; if (typeof content === 'string') return content; return content.map((part) => (typeof part?.text === 'string' ? part.text : '')).join(''); } -/** Normalizes messages into the `{ role, content }` form @axlabs/venice-e2ee expects. */ -export function toPlainMessages( - messages: ReadonlyArray>, -): Array<{ role: string; content: string }> { - return messages.map((message) => ({ - role: String(message.role), - content: flattenContent(message.content), - })); +/** Message passed to `@axlabs/venice-e2ee` `encrypt`, including tool-call structure. */ +export interface PlainMessage { + role: string; + content: string | null; + name?: string; + tool_call_id?: string; + tool_calls?: ToolCallDelta[]; +} + +/** Fields read from a chat message before encryption. */ +interface PlainMessageSource { + role: string; + content?: ChatMessageContent; + name?: string; + tool_call_id?: string; + tool_calls?: ToolCallDelta[]; +} + +/** + * Normalizes messages for encryption. Tool-call metadata is preserved in the + * clear. `content: null` stays `null` so a tool-only assistant turn is not + * flattened into an empty string (which the library would otherwise have to + * treat as text). + */ +export function toPlainMessages(messages: ReadonlyArray): PlainMessage[] { + return messages.map((message) => { + const plain: PlainMessage = { + role: String(message.role), + content: + message.content === null || message.content === undefined + ? null + : flattenContent(message.content), + }; + if (typeof message.name === 'string') plain.name = message.name; + if (typeof message.tool_call_id === 'string') plain.tool_call_id = message.tool_call_id; + if (Array.isArray(message.tool_calls)) plain.tool_calls = message.tool_calls; + return plain; + }); } /** Character counts for an incoming chat payload (plaintext, before E2EE). */ @@ -58,9 +90,9 @@ export interface IncomingMessageCharStats { export interface OutgoingE2eeMessageCharStats { messageCount: number; totalChars: number; - /** Hex ciphertext length for encrypted user/system messages. */ + /** Hex ciphertext length for user/system `content`. */ encryptedChars: number; - /** Plaintext length for assistant and other roles left unchanged. */ + /** `content` length for every other role (ciphertext for non-empty tool results). */ plaintextChars: number; } @@ -97,13 +129,13 @@ export function countIncomingMessageChars( /** Totals character counts in message `content` fields after E2EE encryption. */ export function countOutgoingE2eeMessageChars( - messages: ReadonlyArray<{ role: string; content: string }>, + messages: ReadonlyArray<{ role: string; content: string | null }>, ): OutgoingE2eeMessageCharStats { let encryptedChars = 0; let plaintextChars = 0; for (const message of messages) { - const chars = message.content.length; + const chars = message.content?.length ?? 0; if (message.role === 'user' || message.role === 'system') { encryptedChars += chars; } else { @@ -119,13 +151,51 @@ export function countOutgoingE2eeMessageChars( }; } -/** Builds a full non-streaming chat completion from assembled text. */ +/** + * Merges streaming `delta.tool_calls` into completed calls, keyed by `index`. + * `function.name` and `function.arguments` are concatenated in arrival order. + */ +export function appendToolCallDeltas( + accumulated: AssembledToolCall[], + deltas: readonly ToolCallDelta[], +): void { + for (const delta of deltas) { + const index = typeof delta.index === 'number' ? delta.index : 0; + const current = accumulated[index] ?? { + id: '', + type: 'function', + function: { name: '', arguments: '' }, + }; + if (typeof delta.id === 'string' && delta.id.length > 0) current.id = delta.id; + if (typeof delta.type === 'string' && delta.type.length > 0) current.type = delta.type; + if (typeof delta.function?.name === 'string') current.function.name += delta.function.name; + if (typeof delta.function?.arguments === 'string') { + current.function.arguments += delta.function.arguments; + } + accumulated[index] = current; + } +} + +/** Drops holes left by a sparse tool-call index. */ +export function compactToolCalls(accumulated: AssembledToolCall[]): AssembledToolCall[] { + return accumulated.filter((call) => call != null); +} + +/** Builds a full non-streaming chat completion from assembled text and tool calls. */ export function buildChatCompletion(params: { id: string; model: string; content: string; finishReason?: string; + toolCalls?: AssembledToolCall[]; }): ChatCompletion { + const toolCalls = params.toolCalls ?? []; + const message: ChatCompletion['choices'][number]['message'] = { + role: 'assistant', + content: toolCalls.length > 0 && params.content.length === 0 ? null : params.content, + }; + if (toolCalls.length > 0) message.tool_calls = toolCalls; + return { id: params.id, object: 'chat.completion', @@ -134,7 +204,7 @@ export function buildChatCompletion(params: { choices: [ { index: 0, - message: { role: 'assistant', content: params.content }, + message, finish_reason: params.finishReason ?? 'stop', logprobs: null, }, @@ -187,6 +257,29 @@ export function buildRoleChunk(params: { }; } +/** Builds a streaming chunk carrying `delta.tool_calls`, in the order received. */ +export function buildToolCallsChunk(params: { + id: string; + model: string; + created: number; + toolCalls: ToolCallDelta[]; +}): ChatCompletionChunk { + return { + id: params.id, + object: 'chat.completion.chunk', + created: params.created, + model: params.model, + choices: [ + { + index: 0, + delta: { tool_calls: params.toolCalls }, + finish_reason: null, + logprobs: null, + }, + ], + }; +} + /** Builds the terminating chunk carrying the finish reason. */ export function buildFinalChunk(params: { id: string; diff --git a/src/e2ee/e2ee.service.ts b/src/e2ee/e2ee.service.ts index fcb878e..9c8600d 100644 --- a/src/e2ee/e2ee.service.ts +++ b/src/e2ee/e2ee.service.ts @@ -3,7 +3,9 @@ import { Injectable, Logger, OnModuleInit, ServiceUnavailableException } from '@ import type { createVeniceE2EE, E2EESession, + E2EEMessage, EncryptedPayload, + DecryptedStreamEvent, AttestationResult, AttestationResponse, DcapVerifier, @@ -30,7 +32,7 @@ const importEsm = new Function('specifier', 'return import(specifier);') as ( specifier: string, ) => Promise; -export type { E2EESession, AttestationResult }; +export type { E2EESession, DecryptedStreamEvent, AttestationResult }; /** * Wrapper around the `@axlabs/venice-e2ee` library. Owns the single E2EE instance, @@ -237,17 +239,20 @@ export class E2eeService implements OnModuleInit { this.gpuAttestationCache.set(modelId, now + this.config.sessionTtlMs); } - /** Encrypts messages for an established session. */ - async encrypt( - messages: Array<{ role: string; content: string }>, - session: E2EESession, - ): Promise { + /** Encrypts messages for an established session. Tool metadata is left in the clear. */ + async encrypt(messages: E2EEMessage[], session: E2EESession): Promise { this.logger.debug(`Encrypting ${messages.length} message(s) for "${session.modelId}"`); return this.instance.encrypt(messages, session); } - /** Decrypts a Venice SSE stream, yielding plaintext text deltas. */ - decryptStream(body: ReadableStream, session: E2EESession): AsyncGenerator { + /** + * Decrypts a Venice SSE stream into content, tool-call, and finish events. + * Tool-call deltas are not ciphertext. + */ + decryptStream( + body: ReadableStream, + session: E2EESession, + ): AsyncGenerator { return this.instance.decryptStream(body, session); } diff --git a/test/helpers/test-app.ts b/test/helpers/test-app.ts index bb4bcae..342955a 100644 --- a/test/helpers/test-app.ts +++ b/test/helpers/test-app.ts @@ -1,4 +1,5 @@ import type { INestApplication } from '@nestjs/common'; +import type { DecryptedStreamEvent } from '@axlabs/venice-e2ee'; import { Test } from '@nestjs/testing'; import { vi } from 'vitest'; import { AppModule } from '../../src/app.module'; @@ -36,9 +37,16 @@ export function readableFrom(parts: string[]): ReadableStream { }); } -/** Async generator helper mirroring `@axlabs/venice-e2ee`'s decryptStream output. */ -export async function* yieldTexts(texts: string[]): AsyncGenerator { - for (const text of texts) yield text; +/** Async generator of decrypted stream events, matching `@axlabs/venice-e2ee`. */ +export async function* yieldEvents( + events: DecryptedStreamEvent[], +): AsyncGenerator { + for (const event of events) yield event; +} + +/** Content-only stream events, for tests that do not exercise tool calls. */ +export async function* yieldTexts(texts: string[]): AsyncGenerator { + for (const text of texts) yield { kind: 'content', content: text }; } /** @@ -59,8 +67,14 @@ export async function createTestApp(): Promise { const e2ee: MockedE2ee = { createSession: vi.fn(async (modelId: string) => ({ modelId })), - encrypt: vi.fn(async (messages: Array<{ role: string; content: string }>) => ({ - encryptedMessages: messages.map((m) => ({ role: m.role, content: `ENC(${m.content})` })), + encrypt: vi.fn(async (messages: Array<{ role: string; content?: string | null }>) => ({ + encryptedMessages: messages.map((message) => ({ + ...message, + content: + typeof message.content === 'string' && message.content.length > 0 + ? `ENC(${message.content})` + : (message.content ?? null), + })), headers: { 'X-Venice-TEE-Client-Pub-Key': 'client-pub' }, veniceParameters: { enable_e2ee: true }, })), diff --git a/test/openai-compliance.spec.ts b/test/openai-compliance.spec.ts index b471bfb..987dbd0 100644 --- a/test/openai-compliance.spec.ts +++ b/test/openai-compliance.spec.ts @@ -1,7 +1,13 @@ import type { INestApplication } from '@nestjs/common'; import request from 'supertest'; import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest'; -import { createTestApp, readableFrom, yieldTexts, type TestHarness } from './helpers/test-app'; +import { + createTestApp, + readableFrom, + yieldEvents, + yieldTexts, + type TestHarness, +} from './helpers/test-app'; /** * OpenAI API compliance tests. These verify the proxy speaks the shapes that @@ -159,7 +165,12 @@ describe('OpenAI-compatible proxy', () => { expect(res.body).toMatchObject({ object: 'chat.completion', model: 'e2ee-qwen3-5-122b-a10b', - choices: [{ message: { role: 'assistant', content: 'Buffered answer' } }], + choices: [ + { + message: { role: 'assistant', content: 'Buffered answer' }, + finish_reason: 'stop', + }, + ], }); // Even for a non-streaming client request, upstream is forced to stream. expect(harness.venice.postJson).toHaveBeenCalledWith( @@ -168,6 +179,200 @@ describe('OpenAI-compatible proxy', () => { expect.anything(), ); }); + + it('streams tool-call deltas in order and finishes with tool_calls', async () => { + const first = { + index: 0, + id: 'call_weather', + type: 'function', + function: { name: 'get_weather', arguments: '' }, + }; + const second = { index: 0, function: { arguments: '{"city":"Paris"}' } }; + harness.e2ee.decryptStream.mockReturnValue( + yieldEvents([ + { kind: 'content', content: 'Checking ' }, + { kind: 'content', content: 'weather' }, + { kind: 'tool_calls', tool_calls: [first] }, + { kind: 'tool_calls', tool_calls: [second] }, + { kind: 'finish', finish_reason: 'tool_calls' }, + ]), + ); + harness.venice.postJson.mockResolvedValue({ ok: true, body: {} }); + + const res = await request(app.getHttpServer()) + .post('/v1/chat/completions') + .send({ + model: 'e2ee-qwen3-5-122b-a10b', + messages: [{ role: 'user', content: 'Weather?' }], + stream: true, + }) + .expect(200); + + const chunks = parseSse(res.text); + const deltas = chunks.map((chunk) => chunk.choices[0].delta); + expect(deltas[0]).toMatchObject({ role: 'assistant' }); + expect( + deltas.filter((delta) => typeof delta.content === 'string' && delta.content.length > 0), + ).toEqual([{ content: 'Checking ' }, { content: 'weather' }]); + expect( + deltas.filter((delta) => Array.isArray(delta.tool_calls)).map((delta) => delta.tool_calls), + ).toEqual([[first], [second]]); + expect(chunks[chunks.length - 1].choices[0]).toMatchObject({ + delta: {}, + finish_reason: 'tool_calls', + }); + expect(chunks.some((chunk) => chunk.choices[0].finish_reason === 'stop')).toBe(false); + }); + + it('assembles tool calls by index for clients that do not stream', async () => { + harness.e2ee.decryptStream.mockReturnValue( + yieldEvents([ + { kind: 'content', content: 'Sure' }, + { + kind: 'tool_calls', + tool_calls: [ + { + index: 0, + id: 'call_a', + type: 'function', + function: { name: 'get_weather', arguments: '' }, + }, + { + index: 1, + id: 'call_b', + type: 'function', + function: { name: 'get_time', arguments: '' }, + }, + ], + }, + { + kind: 'tool_calls', + tool_calls: [ + { index: 0, function: { arguments: '{"city":' } }, + { index: 1, function: { arguments: '{"tz":"UTC"}' } }, + ], + }, + { kind: 'tool_calls', tool_calls: [{ index: 0, function: { arguments: '"Paris"}' } }] }, + { kind: 'finish', finish_reason: 'tool_calls' }, + ]), + ); + harness.venice.postJson.mockResolvedValue({ ok: true, body: {} }); + + const res = await request(app.getHttpServer()) + .post('/v1/chat/completions') + .send({ + model: 'e2ee-qwen3-5-122b-a10b', + messages: [{ role: 'user', content: 'Weather and time?' }], + }) + .expect(200); + + expect(res.body.choices[0]).toMatchObject({ + finish_reason: 'tool_calls', + message: { + role: 'assistant', + content: 'Sure', + tool_calls: [ + { + id: 'call_a', + type: 'function', + function: { name: 'get_weather', arguments: '{"city":"Paris"}' }, + }, + { + id: 'call_b', + type: 'function', + function: { name: 'get_time', arguments: '{"tz":"UTC"}' }, + }, + ], + }, + }); + }); + + it('returns null content when the assistant turn is only a tool call', async () => { + harness.e2ee.decryptStream.mockReturnValue( + yieldEvents([ + { + kind: 'tool_calls', + tool_calls: [ + { + index: 0, + id: 'call_only', + type: 'function', + function: { name: 'lookup', arguments: '{}' }, + }, + ], + }, + { kind: 'finish', finish_reason: 'length' }, + ]), + ); + harness.venice.postJson.mockResolvedValue({ ok: true, body: {} }); + + const res = await request(app.getHttpServer()) + .post('/v1/chat/completions') + .send({ + model: 'e2ee-qwen3-5-122b-a10b', + messages: [{ role: 'user', content: 'Look it up' }], + }) + .expect(200); + + expect(res.body.choices[0].finish_reason).toBe('length'); + expect(res.body.choices[0].message).toEqual({ + role: 'assistant', + content: null, + tool_calls: [ + { id: 'call_only', type: 'function', function: { name: 'lookup', arguments: '{}' } }, + ], + }); + }); + + it('accepts a follow-up whose assistant message has null content and tool fields', async () => { + harness.e2ee.decryptStream.mockReturnValue( + yieldEvents([ + { kind: 'content', content: 'It is sunny.' }, + { kind: 'finish', finish_reason: 'stop' }, + ]), + ); + harness.venice.postJson.mockResolvedValue({ ok: true, body: {} }); + + const toolCalls = [ + { + id: 'call_weather', + type: 'function', + function: { name: 'get_weather', arguments: '{"city":"Paris"}' }, + }, + ]; + + const res = await request(app.getHttpServer()) + .post('/v1/chat/completions') + .send({ + model: 'e2ee-qwen3-5-122b-a10b', + messages: [ + { role: 'user', content: 'Weather in Paris?' }, + { role: 'assistant', content: null, tool_calls: toolCalls }, + { + role: 'tool', + tool_call_id: 'call_weather', + name: 'get_weather', + content: '22C and sunny', + }, + ], + }) + .expect(200); + + expect(res.body.choices[0].message.content).toBe('It is sunny.'); + expect(harness.e2ee.encrypt).toHaveBeenCalledWith( + [ + { role: 'user', content: 'Weather in Paris?' }, + { role: 'assistant', content: null, tool_calls: toolCalls }, + { + role: 'tool', + content: '22C and sunny', + name: 'get_weather', + tool_call_id: 'call_weather', + }, + ], + expect.anything(), + ); + }); }); describe('path compatibility', () => { @@ -194,3 +399,17 @@ describe('OpenAI-compatible proxy', () => { }); }); }); + +function parseSse(body: string): Array<{ + choices: Array<{ delta: Record; finish_reason: string | null }>; +}> { + return body + .split('\n\n') + .filter((block) => block.startsWith('data: ') && !block.includes('[DONE]')) + .map( + (block) => + JSON.parse(block.slice('data: '.length)) as { + choices: Array<{ delta: Record; finish_reason: string | null }>; + }, + ); +} diff --git a/test/openai.util.spec.ts b/test/openai.util.spec.ts index 7431ef4..136cf19 100644 --- a/test/openai.util.spec.ts +++ b/test/openai.util.spec.ts @@ -1,9 +1,12 @@ import { describe, expect, it } from 'vitest'; import { + appendToolCallDeltas, buildChatCompletion, buildContentChunk, buildFinalChunk, buildRoleChunk, + buildToolCallsChunk, + compactToolCalls, countIncomingMessageChars, countOutgoingE2eeMessageChars, flattenContent, @@ -13,6 +16,7 @@ import { toPlainMessages, SSE_DONE, } from '../src/common/utils/openai.util'; +import type { AssembledToolCall } from '../src/common/interfaces/openai.interface'; describe('isE2EEModel', () => { it('detects e2ee- prefixed models', () => { @@ -55,6 +59,34 @@ describe('toPlainMessages', () => { { role: 'user', content: 'hi' }, ]); }); + + it('keeps null content and tool fields for a follow-up turn', () => { + const toolCalls = [ + { + id: 'call_weather', + type: 'function', + function: { name: 'get_weather', arguments: '{"city":"Paris"}' }, + }, + ]; + const result = toPlainMessages([ + { role: 'assistant', content: null, tool_calls: toolCalls }, + { + role: 'tool', + content: '22C and sunny', + tool_call_id: 'call_weather', + name: 'get_weather', + }, + ]); + expect(result).toEqual([ + { role: 'assistant', content: null, tool_calls: toolCalls }, + { + role: 'tool', + content: '22C and sunny', + tool_call_id: 'call_weather', + name: 'get_weather', + }, + ]); + }); }); describe('message char stats', () => { @@ -77,13 +109,14 @@ describe('message char stats', () => { it('counts outgoing E2EE content as encrypted vs plaintext', () => { const stats = countOutgoingE2eeMessageChars([ { role: 'user', content: 'a'.repeat(186) }, - { role: 'assistant', content: 'plain reply' }, + { role: 'assistant', content: null }, + { role: 'tool', content: 'cipher' }, ]); expect(stats).toEqual({ - messageCount: 2, - totalChars: 186 + 11, + messageCount: 3, + totalChars: 186 + 6, encryptedChars: 186, - plaintextChars: 11, + plaintextChars: 6, }); }); }); @@ -110,8 +143,50 @@ describe('OpenAI builders', () => { expect(content.choices[0].delta.content).toBe('tok'); expect(content.choices[0].finish_reason).toBeNull(); - const final = buildFinalChunk({ id: 'c1', model: 'm', created }); - expect(final.choices[0].finish_reason).toBe('stop'); + const final = buildFinalChunk({ id: 'c1', model: 'm', created, finishReason: 'tool_calls' }); + expect(final.choices[0].finish_reason).toBe('tool_calls'); + expect(final.choices[0].delta).toEqual({}); + + const call = { + index: 0, + id: 'call_1', + type: 'function', + function: { name: 'lookup', arguments: '{}' }, + }; + const toolChunk = buildToolCallsChunk({ id: 'c1', model: 'm', created, toolCalls: [call] }); + expect(toolChunk.choices[0].delta.tool_calls).toEqual([call]); + expect(toolChunk.choices[0].finish_reason).toBeNull(); + }); + + it('assembles tool-call deltas by index', () => { + const accumulated: AssembledToolCall[] = []; + appendToolCallDeltas(accumulated, [ + { + index: 0, + id: 'call_a', + type: 'function', + function: { name: 'get_weather', arguments: '' }, + }, + { index: 1, id: 'call_b', type: 'function', function: { name: 'get_time', arguments: '' } }, + ]); + appendToolCallDeltas(accumulated, [ + { index: 0, function: { arguments: '{"city":' } }, + { index: 1, function: { arguments: '{"tz":"UTC"}' } }, + ]); + appendToolCallDeltas(accumulated, [{ index: 0, function: { arguments: '"Paris"}' } }]); + + expect(compactToolCalls(accumulated)).toEqual([ + { + id: 'call_a', + type: 'function', + function: { name: 'get_weather', arguments: '{"city":"Paris"}' }, + }, + { + id: 'call_b', + type: 'function', + function: { name: 'get_time', arguments: '{"tz":"UTC"}' }, + }, + ]); }); }); diff --git a/test/tool-encryption.spec.ts b/test/tool-encryption.spec.ts new file mode 100644 index 0000000..d9d6281 --- /dev/null +++ b/test/tool-encryption.spec.ts @@ -0,0 +1,87 @@ +import { describe, expect, it } from 'vitest'; +import { + createVeniceE2EE, + deriveAESKey, + fromHex, + generateKeypair, + toHex, + type E2EESession, +} from '@axlabs/venice-e2ee'; + +async function sessionPair(): Promise<{ session: E2EESession; modelPrivateKey: Uint8Array }> { + const client = generateKeypair(); + const model = generateKeypair(); + const aesKey = await deriveAESKey(client.privateKey, model.pubKeyHex); + return { + modelPrivateKey: model.privateKey, + session: { + privateKey: client.privateKey, + publicKey: client.publicKey, + pubKeyHex: client.pubKeyHex, + modelPubKeyHex: model.pubKeyHex, + aesKey, + modelId: 'e2ee-test', + created: Date.now(), + }, + }; +} + +async function decryptWire(cipherHex: string, modelPrivateKey: Uint8Array): Promise { + const raw = fromHex(cipherHex); + const clientPubHex = toHex(raw.slice(0, 65)); + const iv = raw.slice(65, 77); + const ciphertext = raw.slice(77); + const aesKey = await deriveAESKey(modelPrivateKey, clientPubHex); + const plaintext = await crypto.subtle.decrypt({ name: 'AES-GCM', iv }, aesKey, ciphertext); + return new TextDecoder().decode(plaintext); +} + +describe('tool-result encryption', () => { + it('encrypts tool-result content and keeps call metadata in the clear', async () => { + const { session, modelPrivateKey } = await sessionPair(); + const e2ee = createVeniceE2EE({ apiKey: 'test', verifyAttestation: false }); + const secret = '22C and sunny'; + + const { encryptedMessages } = await e2ee.encrypt( + [ + { + role: 'assistant', + content: null, + tool_calls: [ + { + id: 'call_weather', + type: 'function', + function: { name: 'get_weather', arguments: '{"city":"Paris"}' }, + }, + ], + }, + { + role: 'tool', + name: 'get_weather', + tool_call_id: 'call_weather', + content: secret, + }, + ], + session, + ); + + expect(encryptedMessages[0]).toEqual({ + role: 'assistant', + content: null, + tool_calls: [ + { + id: 'call_weather', + type: 'function', + function: { name: 'get_weather', arguments: '{"city":"Paris"}' }, + }, + ], + }); + + const toolMessage = encryptedMessages[1]; + expect(toolMessage.name).toBe('get_weather'); + expect(toolMessage.tool_call_id).toBe('call_weather'); + expect(toolMessage.content).not.toBe(secret); + expect(toolMessage.content).toMatch(/^[0-9a-f]+$/i); + expect(await decryptWire(toolMessage.content!, modelPrivateKey)).toBe(secret); + }); +});