|
|
|
@@ -0,0 +1,220 @@
|
|
|
|
|
#!/usr/bin/env bun
|
|
|
|
|
/**
|
|
|
|
|
* Zesdex WebSocket interface — real-time bidirectional communication.
|
|
|
|
|
* Mirrors `apps/interfaces/ws/src/lib.rs`.
|
|
|
|
|
*
|
|
|
|
|
* Security: accepts an optional `?token=` query param. When `ZESDEX_WS_TOKEN`
|
|
|
|
|
* env is set, connections MUST present a matching token, else rejected —
|
|
|
|
|
* prevents the endpoint from being used as an open LLM proxy.
|
|
|
|
|
*
|
|
|
|
|
* Protocol (JSON text frames):
|
|
|
|
|
* client → { "type": "prompt", "message": "...", "model"?: "..." }
|
|
|
|
|
* server → { "type": "connected", "session": ..., "message": "..." } (on connect)
|
|
|
|
|
* server → { "type": "token", "content": "..." } (streaming)
|
|
|
|
|
* server → { "type": "done" } | { "type": "error", "message": "..." } (terminal)
|
|
|
|
|
* server → { "type": "echo", "data": "..." } (fallback)
|
|
|
|
|
*/
|
|
|
|
|
import {
|
|
|
|
|
JsonSettingsRepository,
|
|
|
|
|
JsonAppConfigRepository,
|
|
|
|
|
LlmClient,
|
|
|
|
|
resolveApiKey,
|
|
|
|
|
InfrastructureToolExecutor,
|
|
|
|
|
allTools,
|
|
|
|
|
toolDefs,
|
|
|
|
|
} from "@zesdex/infrastructure";
|
|
|
|
|
import type { ToolCtx } from "@zesdex/infrastructure";
|
|
|
|
|
import { AgentTurnServiceImpl } from "@zesdex/application";
|
|
|
|
|
import {
|
|
|
|
|
newStore,
|
|
|
|
|
ensureStoreDirs,
|
|
|
|
|
type TurnEvent,
|
|
|
|
|
type TurnEventSink,
|
|
|
|
|
type AgentTurnParams,
|
|
|
|
|
userMessage,
|
|
|
|
|
resolveEffectiveModel,
|
|
|
|
|
} from "@zesdex/domain";
|
|
|
|
|
|
|
|
|
|
/** Shared application state for the WS server. */
|
|
|
|
|
export interface WsState {
|
|
|
|
|
store_base_dir: string;
|
|
|
|
|
session_id: string | null;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/** Minimal send interface used by streaming helpers. */
|
|
|
|
|
export interface SocketSend {
|
|
|
|
|
send(json: string): void;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/** Array-backed turn event sink (push + drain). */
|
|
|
|
|
class ArraySink implements TurnEventSink {
|
|
|
|
|
constructor(public events: TurnEvent[] = []) {}
|
|
|
|
|
push(event: TurnEvent): void {
|
|
|
|
|
this.events.push(event);
|
|
|
|
|
}
|
|
|
|
|
drain(): TurnEvent[] {
|
|
|
|
|
const out = this.events;
|
|
|
|
|
this.events = [];
|
|
|
|
|
return out;
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/** Validate the token against ZESDEX_WS_TOKEN (when configured). */
|
|
|
|
|
function tokenAllowed(url: URL): boolean {
|
|
|
|
|
const configured = process.env.ZESDEX_WS_TOKEN;
|
|
|
|
|
const t = url.searchParams.get("token");
|
|
|
|
|
if (configured) {
|
|
|
|
|
return t === configured;
|
|
|
|
|
}
|
|
|
|
|
return true;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/** Wire a turn service for one prompt (matches the CLI composition root). */
|
|
|
|
|
async function buildSingleShot(
|
|
|
|
|
sessionDir: string,
|
|
|
|
|
eventSink: TurnEventSink,
|
|
|
|
|
): Promise<{ turnService: AgentTurnServiceImpl; params: AgentTurnParams }> {
|
|
|
|
|
const store = newStore();
|
|
|
|
|
await ensureStoreDirs(store);
|
|
|
|
|
|
|
|
|
|
const settingsRepo = new JsonSettingsRepository();
|
|
|
|
|
const appConfigRepo = new JsonAppConfigRepository();
|
|
|
|
|
const settings = await settingsRepo.load(store.base_dir);
|
|
|
|
|
const appConfig = await appConfigRepo.load(store.base_dir);
|
|
|
|
|
|
|
|
|
|
const provider = settings.provider;
|
|
|
|
|
const model = resolveEffectiveModel(settings, appConfig);
|
|
|
|
|
const baseUrl = appConfig.providers[provider]?.api_base ?? undefined;
|
|
|
|
|
const apiKey = resolveApiKey(settings, appConfig);
|
|
|
|
|
|
|
|
|
|
const llmClient = new LlmClient(apiKey, model, baseUrl);
|
|
|
|
|
const toolCtx = {
|
|
|
|
|
sessionDir,
|
|
|
|
|
workspaces: [sessionDir],
|
|
|
|
|
turnEvents: eventSink,
|
|
|
|
|
workflowFindings: [],
|
|
|
|
|
} as unknown as ToolCtx;
|
|
|
|
|
const executor = new InfrastructureToolExecutor(toolCtx);
|
|
|
|
|
const turnService = new AgentTurnServiceImpl(llmClient, executor, toolDefs(allTools()) as never);
|
|
|
|
|
|
|
|
|
|
const params: AgentTurnParams = {
|
|
|
|
|
messages: [userMessage("")],
|
|
|
|
|
session_dir: sessionDir,
|
|
|
|
|
workspace_roots: [sessionDir],
|
|
|
|
|
turn_events: eventSink,
|
|
|
|
|
in_flight: { value: false },
|
|
|
|
|
abort: new AbortController(),
|
|
|
|
|
api_key: apiKey,
|
|
|
|
|
model,
|
|
|
|
|
api_base: baseUrl,
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
return { turnService, params };
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Start the WS server. `onPrompt` is invoked with the client's prompt payload,
|
|
|
|
|
* and `forward` sends raw strings to the client.
|
|
|
|
|
*/
|
|
|
|
|
export function startWsServer(port: number): { stop: () => void } {
|
|
|
|
|
const state: WsState = { store_base_dir: ".", session_id: null };
|
|
|
|
|
|
|
|
|
|
const server = Bun.serve({
|
|
|
|
|
port,
|
|
|
|
|
fetch(req, server) {
|
|
|
|
|
const url = new URL(req.url);
|
|
|
|
|
if (url.pathname === "/ws") {
|
|
|
|
|
if (!tokenAllowed(url)) {
|
|
|
|
|
return new Response("missing or invalid token", { status: 401 });
|
|
|
|
|
}
|
|
|
|
|
if (server.upgrade(req, {})) {
|
|
|
|
|
return undefined;
|
|
|
|
|
}
|
|
|
|
|
return new Response("upgrade failed", { status: 400 });
|
|
|
|
|
}
|
|
|
|
|
return new Response("Not found", { status: 404 });
|
|
|
|
|
},
|
|
|
|
|
websocket: {
|
|
|
|
|
open(ws) {
|
|
|
|
|
const sender: SocketSend = { send: (s) => ws.send(s) };
|
|
|
|
|
// Welcome message on connect.
|
|
|
|
|
sender.send(
|
|
|
|
|
JSON.stringify({
|
|
|
|
|
type: "connected",
|
|
|
|
|
session: state.session_id,
|
|
|
|
|
message: "Connected to Zesdex WebSocket server",
|
|
|
|
|
}),
|
|
|
|
|
);
|
|
|
|
|
},
|
|
|
|
|
message(ws, msg) {
|
|
|
|
|
const text = typeof msg === "string" ? msg : JSON.stringify(msg);
|
|
|
|
|
if (!text.startsWith("{")) {
|
|
|
|
|
ws.send(JSON.stringify({ type: "echo", data: text }));
|
|
|
|
|
return;
|
|
|
|
|
}
|
|
|
|
|
let val: Record<string, unknown>;
|
|
|
|
|
try {
|
|
|
|
|
val = JSON.parse(text) as Record<string, unknown>;
|
|
|
|
|
} catch {
|
|
|
|
|
ws.send(JSON.stringify({ type: "echo", data: text }));
|
|
|
|
|
return;
|
|
|
|
|
}
|
|
|
|
|
if (val.type === "prompt" && typeof val.message === "string") {
|
|
|
|
|
const sender: SocketSend = { send: (s) => ws.send(s) };
|
|
|
|
|
void runPrompt(sender, val.message, typeof val.model === "string" ? val.model : undefined);
|
|
|
|
|
} else {
|
|
|
|
|
ws.send(JSON.stringify({ type: "echo", data: text }));
|
|
|
|
|
}
|
|
|
|
|
},
|
|
|
|
|
close() {},
|
|
|
|
|
},
|
|
|
|
|
});
|
|
|
|
|
console.log(`zesdex-ws listening on ws://0.0.0.0:${server.port}`);
|
|
|
|
|
return { stop: () => server.stop(true) };
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/** Run a single prompt and stream events to the client. */
|
|
|
|
|
async function runPrompt(ws: SocketSend, message: string, modelOverride?: string): Promise<void> {
|
|
|
|
|
const sessionDir = process.cwd();
|
|
|
|
|
const sink = new ArraySink();
|
|
|
|
|
try {
|
|
|
|
|
const { turnService, params } = await buildSingleShot(sessionDir, sink);
|
|
|
|
|
if (modelOverride) params.model = modelOverride;
|
|
|
|
|
params.messages = [userMessage(message)];
|
|
|
|
|
|
|
|
|
|
// Poll the sink and forward events to the client.
|
|
|
|
|
const poll = setInterval(() => {
|
|
|
|
|
const evs = sink.drain();
|
|
|
|
|
for (const ev of evs) {
|
|
|
|
|
if (ev.kind === "stream_token") {
|
|
|
|
|
ws.send(JSON.stringify({ type: "token", content: ev.content }));
|
|
|
|
|
} else if (ev.kind === "done") {
|
|
|
|
|
ws.send(JSON.stringify({ type: "done" }));
|
|
|
|
|
clearInterval(poll);
|
|
|
|
|
} else if (ev.kind === "error") {
|
|
|
|
|
ws.send(JSON.stringify({ type: "error", message: ev.message }));
|
|
|
|
|
clearInterval(poll);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}, 50);
|
|
|
|
|
|
|
|
|
|
try {
|
|
|
|
|
await turnService.runTurn(params);
|
|
|
|
|
// Ensure a terminal done/error is sent even if none auto-emitted
|
|
|
|
|
// while the poll stopped early.
|
|
|
|
|
const remaining = sink.drain();
|
|
|
|
|
for (const ev of remaining) {
|
|
|
|
|
if (ev.kind === "stream_token") ws.send(JSON.stringify({ type: "token", content: ev.content }));
|
|
|
|
|
else if (ev.kind === "done") ws.send(JSON.stringify({ type: "done" }));
|
|
|
|
|
else if (ev.kind === "error") ws.send(JSON.stringify({ type: "error", message: ev.message }));
|
|
|
|
|
}
|
|
|
|
|
if (!remaining.some((e) => e.kind === "done" || e.kind === "error")) {
|
|
|
|
|
ws.send(JSON.stringify({ type: "done" }));
|
|
|
|
|
}
|
|
|
|
|
} finally {
|
|
|
|
|
clearInterval(poll);
|
|
|
|
|
}
|
|
|
|
|
} catch (e) {
|
|
|
|
|
ws.send(JSON.stringify({ type: "error", message: (e as Error).message }));
|
|
|
|
|
}
|
|
|
|
|
}
|
asepharyana