feat(ai): extend the injection guard to the stream - #668
Conversation
a08550b to
18f7aa1
Compare
18f7aa1 to
571263c
Compare
571263c to
e64c1a7
Compare
There was a problem hiding this comment.
Pull request overview
Extends the existing prompt-injection / marker-leak defense from one-shot completions to streamed AI suggestions by introducing nonce-based spotlighting, deterministic leak detection, and a streaming “holdback” guard that prevents nonce substrings from crossing chunk boundaries.
Changes:
- Add spotlighting markers + per-request random nonce in the prompt, and reject outputs that reproduce the nonce.
- Add streaming leak defense via a holdback guard wired into
streamText(...).experimental_transform. - Add an Express
/integration/ai/streamroute that proxies the UI message stream response; expand tests and shared test helpers accordingly.
Reviewed changes
Copilot reviewed 20 out of 20 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| test/services/askAi.test.ts | Expands AIService tests to cover spotlighting, rejection reporting, and stream guard wiring. |
| test/services/askAi-spotlighting.test.ts | Adds unit coverage for prompt marker wrapping and nonce generation/collision handling. |
| test/services/askAi-leak-detector.test.ts | Adds unit tests for deterministic nonce leak detection behavior. |
| test/services/askAi-holdback.test.ts | Adds unit tests for the streaming holdback guard behavior and edge cases. |
| test/integrations/vercel-ai.test.ts | Adds tests for streamText usage and the guarded transform wiring/semantics. |
| test/integrations/github-routes.test.ts | Refactors to reuse the new Express request helper. |
| test/integrations/ai-routes.test.ts | Adds route-level tests for the new SSE proxy endpoint and auth/validation paths. |
| test/helpers/expressRequest.ts | Introduces a reusable Express request/response harness that supports pipe() streaming. |
| src/services/types.ts | Exports the Event type for reuse (e.g., AI service internals). |
| src/services/askAi/security/spotlighting.ts | Implements nonce-based marker wrapping (buildEventPrompt) + system spotlighting instruction. |
| src/services/askAi/security/leakDetector.ts | Adds isLeaked + fallback message constant for suggestion rejection. |
| src/services/askAi/security/holdback.ts | Implements streaming holdback guard (createLeakGuard) to detect nonce across deltas. |
| src/services/askAi/inputs/eventSolving.ts | Documents that this serialization is unwrapped/untrusted and must be wrapped via spotlighting. |
| src/services/ai.ts | Updates AIService to use spotlighting + leak detection, and adds a streaming suggestion method. |
| src/integrations/vercel-ai/routes.ts | Adds /integration/ai/stream proxy route for the UI message stream response. |
| src/integrations/vercel-ai/index.ts | Adds vercelAIApi.stream and the guarded transform that applies the stream guard to deltas. |
| src/index.ts | Registers the new AI assistant routes on the Express app. |
| src/directives/requireUserInWorkspace.ts | Exports checkUserInWorkspaceByProjectId for reuse in the new route. |
| package.json | Bumps package version. |
| .eslintrc.js | Declares Fetch/Streams globals used in tests and streaming code. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
The stream route mapped any error from streamSuggestion to a 404 "Event not found", including failures unrelated to the event lookup (e.g. a stream construction error). Only the exact "Event not found" error is now reported as 404; anything else is forwarded to Express's error handling. Flagged by Copilot while reviewing #668, against code this PR added.
e64c1a7 to
983fded
Compare
78359fa to
6b2ec10
Compare
6b2ec10 to
acda192
Compare
acda192 to
16b0278
Compare
16b0278 to
a480368
Compare
0bcb063 to
00b1962
Compare
New HTTP route GET /integration/ai/stream added. This route calls Ask AI service about specified event and responds text/event-stream. Reponse carrying AiStream which represent AiStreamPart sequence: either text-delta (text-increments generated by AI assistant) or error (failure description during response generation). NOTE: Response doesn't carry reasoning, tooling and start/end parts since they're not required yet. Route checks workspace membership before calling Ask AI. For this purpose function checkUserInWorkspaceByProjectId became exported. Failed membership check leads to response with 403. Also route checks if specified event exists. Failed check leads to response with 404. Aborting request cancel Ask AI suggestion generation. For this purpose AbortController is declared as an eslint global: it is on globalThis since Node 15, but eslint's node env predates it.
The prompt-injection check runs once on a finished answer, but a streamed answer leaves the server as it is written, so a nonce split across two deltas passes a per-delta check untouched. Rejecting a stream also left the model running, since nothing told the transport to stop generating text nobody would see.
00b1962 to
dadfeef
Compare
There was a problem hiding this comment.
🟡 Changes recommended
The critical error-handling path can expose unscannable model output, including the nonce, to clients.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
src/services/askAi/security/holdback.ts:133
- The sentence is missing an article; use “with a guard.”
* Wrap a suggestion stream with guard.
- Files reviewed: 7/8 changed files
- Comments generated: 2
- Review effort level: Balanced
| /** Report unhandled part loudly */ | ||
| const unscannable: never = part; | ||
|
|
||
| throw new Error(`Unscannable suggestion part: ${JSON.stringify(unscannable)}`); |
| } ]); | ||
| }); | ||
|
|
||
| it('should report a rejection once nounce echoes', async () => { |
The one-shot path scans the finished answer before returning it. A streamed answer leaves the server while it is still being written, so a nonce split across two deltas would pass a per-delta check and reach the client.
The stream now runs through a guard that withholds the last
nonce.length - 1characters and scans them together with each new delta, releasing only text that can no longer begin the nonce. With the nonce shortened toa3f19cfor legibility, where the real one is 32 hex characters:That length is the exact minimum: an occurrence spans
nonce.lengthcharacters, so holding one less leaves it inside a single scanned window. Whatever is still held goes out when the answer ends.Detection is unchanged and shared with the one-shot path. The guard wraps the stream in
AskAiService, which is where the nonce is, and the wrapped stream is the only onestreamSuggestionreturns. A part the guard cannot scan raises rather than being forwarded, so a variant added toSuggestionPartlater cannot pass unchecked.A rejection goes out as an
errorpart and not as another text delta: the client concatenates deltas, so a fallback message sent that way would land glued to the truncated prefix already on screen. The prefix itself cannot be taken back, and the holdback cuts it mid-word.