Add streaming to Engine orchestrator (#6094)

# Description of Changes
Adds a streaming endpoint to the Java AI orchestrator
(`/api/v1/ai/orchestrate/stream` in addition to the existing
`/api/v1/ai/orchestrate`). This allows the caller to get updates of what
stage of orchestration is being run at the time so UIs can give the user
feedback.

Also contains some dubious Gradle changes to suppress errors coming from
Spotless, when it crashes in Google stuff. I'm not sure if that's
appropriate to add, feel free to ask for changes in review.
This commit is contained in:
James Brunton
2026-04-17 10:01:08 +01:00
committed by GitHub
parent 97ca85d878
commit 688f7f2013
13 changed files with 326 additions and 40 deletions
@@ -2079,6 +2079,13 @@ keywords = "Keywords: odd, even"
numbers = "Numbers/ranges: 5, 10-20"
progressions = "Progressions: 3n, 4n+1"
[chat.progress]
thinking = "Thinking..."
analyzing = "Analysing your request..."
calling_engine = "AI is thinking..."
extracting_content = "Extracting content from your documents..."
processing = "Processing extracted content..."
[certSign]
allSigned = "All participants have signed. Ready to finalize."
awaitingSignatures = "Awaiting signatures"
@@ -12,3 +12,8 @@ export function setupApiInterceptors(client: AxiosInstance): void {
(error) => Promise.reject(error),
);
}
/** Auth headers for raw fetch() calls (SSE streams, etc.). Proprietary overrides with JWT + XSRF. */
export function getAuthHeaders(): Record<string, string> {
return {};
}
@@ -95,20 +95,29 @@ async function refreshAuthToken(client: AxiosInstance): Promise<string> {
}
}
/** Auth headers for raw fetch() calls (SSE streams, etc.). */
export function getAuthHeaders(): Record<string, string> {
const headers: Record<string, string> = {};
const jwt = getJwtTokenFromStorage();
if (jwt) {
headers["Authorization"] = `Bearer ${jwt}`;
}
const xsrf = getXsrfToken();
if (xsrf) {
headers["X-XSRF-TOKEN"] = xsrf;
}
return headers;
}
export function setupApiInterceptors(client: AxiosInstance): void {
// Install request interceptor to add JWT token
client.interceptors.request.use(
(config) => {
const jwtToken = getJwtTokenFromStorage();
const xsrfToken = getXsrfToken();
if (jwtToken && !config.headers.Authorization) {
config.headers.Authorization = `Bearer ${jwtToken}`;
console.debug("[API Client] Added JWT token from localStorage to Authorization header");
}
if (xsrfToken && !config.headers["X-XSRF-TOKEN"]) {
config.headers["X-XSRF-TOKEN"] = xsrfToken;
const authHeaders = getAuthHeaders();
for (const [key, value] of Object.entries(authHeaders)) {
if (!config.headers[key]) {
config.headers[key] = value;
}
}
return config;
@@ -1,5 +1,6 @@
import { createContext, useContext, useReducer, useCallback, type ReactNode } from "react";
import { createContext, useContext, useReducer, useCallback, useRef, type ReactNode } from "react";
import { useAllFiles } from "@app/contexts/FileContext";
import { getAuthHeaders } from "@app/services/apiClientSetup";
export interface ChatMessage {
id: string;
@@ -8,6 +9,13 @@ export interface ChatMessage {
timestamp: number;
}
export enum AiWorkflowPhase {
ANALYZING = "analyzing",
CALLING_ENGINE = "calling_engine",
EXTRACTING_CONTENT = "extracting_content",
PROCESSING = "processing",
}
type AiWorkflowOutcome =
| "answer"
| "not_found"
@@ -37,11 +45,13 @@ interface ChatState {
messages: ChatMessage[];
isOpen: boolean;
isLoading: boolean;
progressPhase: AiWorkflowPhase | null;
}
type ChatAction =
| { type: "ADD_MESSAGE"; message: ChatMessage }
| { type: "SET_LOADING"; loading: boolean }
| { type: "SET_PROGRESS"; phase: AiWorkflowPhase | null }
| { type: "TOGGLE_OPEN" }
| { type: "SET_OPEN"; open: boolean };
@@ -51,6 +61,8 @@ function chatReducer(state: ChatState, action: ChatAction): ChatState {
return { ...state, messages: [...state.messages, action.message] };
case "SET_LOADING":
return { ...state, isLoading: action.loading };
case "SET_PROGRESS":
return { ...state, progressPhase: action.phase };
case "TOGGLE_OPEN":
return { ...state, isOpen: !state.isOpen };
case "SET_OPEN":
@@ -85,10 +97,67 @@ function formatWorkflowResponse(data: AiWorkflowResponse): string {
}
}
/**
* Parses an SSE text stream and invokes callbacks for each named event.
*/
async function consumeSSEStream(
response: Response,
handlers: {
onProgress: (data: { phase: string; timestamp: number }) => void;
onResult: (data: AiWorkflowResponse) => void;
onError: (data: { message: string }) => void;
},
) {
const reader = response.body!.getReader();
const decoder = new TextDecoder();
let buffer = "";
let currentEvent = "";
for (;;) {
const { done, value } = await reader.read();
if (done) break;
buffer += decoder.decode(value, { stream: true });
// SSE frames are separated by double newlines
let boundary = buffer.indexOf("\n\n");
while (boundary !== -1) {
const frame = buffer.slice(0, boundary);
buffer = buffer.slice(boundary + 2);
let dataPayload = "";
for (const line of frame.split("\n")) {
if (line.startsWith("event:")) {
currentEvent = line.slice(6).trim();
} else if (line.startsWith("data:")) {
dataPayload += line.slice(5);
}
}
if (dataPayload) {
try {
const parsed = JSON.parse(dataPayload);
if (currentEvent === "progress") {
handlers.onProgress(parsed);
} else if (currentEvent === "result") {
handlers.onResult(parsed);
} else if (currentEvent === "error") {
handlers.onError(parsed);
}
} catch {
// Skip malformed JSON frames
}
}
currentEvent = "";
boundary = buffer.indexOf("\n\n");
}
}
}
interface ChatContextValue {
messages: ChatMessage[];
isOpen: boolean;
isLoading: boolean;
progressPhase: AiWorkflowPhase | null;
toggleOpen: () => void;
setOpen: (open: boolean) => void;
sendMessage: (content: string) => Promise<void>;
@@ -100,17 +169,24 @@ const initialState: ChatState = {
messages: [],
isOpen: false,
isLoading: false,
progressPhase: null,
};
export function ChatProvider({ children }: { children: ReactNode }) {
const [state, dispatch] = useReducer(chatReducer, initialState);
const { files: activeFiles } = useAllFiles();
const abortRef = useRef<AbortController | null>(null);
const toggleOpen = useCallback(() => dispatch({ type: "TOGGLE_OPEN" }), []);
const setOpen = useCallback((open: boolean) => dispatch({ type: "SET_OPEN", open }), []);
const sendMessage = useCallback(
async (content: string) => {
// Abort any in-flight request
abortRef.current?.abort();
const controller = new AbortController();
abortRef.current = controller;
const userMessage: ChatMessage = {
id: crypto.randomUUID(),
role: "user",
@@ -119,6 +195,7 @@ export function ChatProvider({ children }: { children: ReactNode }) {
};
dispatch({ type: "ADD_MESSAGE", message: userMessage });
dispatch({ type: "SET_LOADING", loading: true });
dispatch({ type: "SET_PROGRESS", phase: null });
try {
const formData = new FormData();
@@ -127,34 +204,73 @@ export function ChatProvider({ children }: { children: ReactNode }) {
formData.append(`fileInputs[${i}].fileInput`, file);
});
const response = await fetch("/api/v1/ai/orchestrate", {
const response = await fetch("/api/v1/ai/orchestrate/stream", {
method: "POST",
body: formData,
headers: getAuthHeaders(),
credentials: "include",
signal: controller.signal,
});
if (!response.ok) {
throw new Error(`AI engine request failed: ${response.status}`);
}
const data: AiWorkflowResponse = await response.json();
const replyContent = formatWorkflowResponse(data);
const assistantMessage: ChatMessage = {
id: crypto.randomUUID(),
role: "assistant",
content: replyContent,
timestamp: Date.now(),
};
dispatch({ type: "ADD_MESSAGE", message: assistantMessage });
} catch {
const errorMessage: ChatMessage = {
id: crypto.randomUUID(),
role: "assistant",
content: "Failed to get a response. The AI engine may not be available yet.",
timestamp: Date.now(),
};
dispatch({ type: "ADD_MESSAGE", message: errorMessage });
let receivedResult = false;
await consumeSSEStream(response, {
onProgress: (data) => {
dispatch({ type: "SET_PROGRESS", phase: data.phase as AiWorkflowPhase });
},
onResult: (data) => {
receivedResult = true;
dispatch({ type: "SET_PROGRESS", phase: null });
const replyContent = formatWorkflowResponse(data);
dispatch({
type: "ADD_MESSAGE",
message: {
id: crypto.randomUUID(),
role: "assistant",
content: replyContent,
timestamp: Date.now(),
},
});
},
onError: (data) => {
receivedResult = true;
dispatch({ type: "SET_PROGRESS", phase: null });
dispatch({
type: "ADD_MESSAGE",
message: {
id: crypto.randomUUID(),
role: "assistant",
content: data.message || "Something went wrong.",
timestamp: Date.now(),
},
});
},
});
if (!receivedResult) {
throw new Error("Stream ended without a result");
}
} catch (e) {
if ((e as Error).name === "AbortError") return;
dispatch({ type: "SET_PROGRESS", phase: null });
dispatch({
type: "ADD_MESSAGE",
message: {
id: crypto.randomUUID(),
role: "assistant",
content: "Failed to get a response. The AI engine may not be available yet.",
timestamp: Date.now(),
},
});
} finally {
dispatch({ type: "SET_LOADING", loading: false });
if (abortRef.current === controller) {
abortRef.current = null;
}
}
},
[activeFiles],
@@ -166,6 +282,7 @@ export function ChatProvider({ children }: { children: ReactNode }) {
messages: state.messages,
isOpen: state.isOpen,
isLoading: state.isLoading,
progressPhase: state.progressPhase,
toggleOpen,
setOpen,
sendMessage,
@@ -1,5 +1,6 @@
import { useRef, useEffect, useState, type KeyboardEvent } from "react";
import { ActionIcon, ScrollArea, TextInput, Stack, Text, Paper, Box, Transition } from "@mantine/core";
import { useTranslation } from "react-i18next";
import { ActionIcon, ScrollArea, TextInput, Stack, Text, Paper, Box, Transition, Loader, Group } from "@mantine/core";
import SendIcon from "@mui/icons-material/Send";
import ChatBubbleOutlineIcon from "@mui/icons-material/ChatBubbleOutline";
import CloseIcon from "@mui/icons-material/Close";
@@ -19,7 +20,8 @@ function ChatMessageBubble({ role, content }: { role: "user" | "assistant"; cont
}
export function ChatPanel() {
const { messages, isOpen, isLoading, toggleOpen, sendMessage } = useChat();
const { t } = useTranslation();
const { messages, isOpen, isLoading, progressPhase, toggleOpen, sendMessage } = useChat();
const [input, setInput] = useState("");
const scrollRef = useRef<HTMLDivElement>(null);
const inputRef = useRef<HTMLInputElement>(null);
@@ -95,9 +97,12 @@ export function ChatPanel() {
{isLoading && (
<div className="chat-message chat-message-assistant">
<Paper className="chat-bubble chat-bubble-assistant" p="xs" radius="md">
<Text size="sm" c="dimmed">
Thinking...
</Text>
<Group gap="xs" wrap="nowrap">
<Loader size="xs" type="dots" />
<Text size="sm" c="dimmed">
{progressPhase ? t(`chat.progress.${progressPhase}`) : t("chat.progress.thinking")}
</Text>
</Group>
</Paper>
</div>
)}