Skip to content

Commit

Permalink
Websocket web (#1005)
Browse files Browse the repository at this point in the history 
* towards durable shared ws connection to server

* 🔨 refactor: streamline WebSocketClient and remove redundancy

* feat: ✨ add progress dispatching to multiple clients

* feat: ✨ add progress reporting for prompt script traces

* feat: ✨ enhance trace handling and add view mode support

* fix react refresh issues

* feat: 🚀 improve run client handling and result updates

* typoings

* feat: ✨ add CORS support, webview, and scrollable Markdown

* add webview

* towards webview
  • Loading branch information
@pelikhan
pelikhan authored Jan 14, 2025
1 parent 0420514 commit c82caf5
Show file tree
Hide file tree
Showing 21 changed files with 765 additions and 588 deletions.
1 change: 1 addition & 0 deletions docs/src/content/docs/reference/cli/commands.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -399,6 +399,7 @@ Options:
--remote-branch <string> Branch to serve from the remote
--remote-force Force pull from remote repository
--remote-install Install dependencies from remote repository
--dispatch-progress Dispatch progress events to all clients
-h, --help display help for command
```

Expand Down
2 changes: 1 addition & 1 deletion package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@
"cache:clear": "cd packages/sample/ && yarn cache:clear",
"test:scripts": "cd packages/sample/ && yarn test:scripts",
"test:scripts:view": "cd packages/sample/ && yarn test:scripts:view",
"serve:cli": "node --watch --watch-path=packages/cli/built packages/cli/built/genaiscript.cjs serve",
"serve:cli": "node --watch --watch-path=packages/cli/built packages/cli/built/genaiscript.cjs serve --dispatch-progress",
"serve:web": "yarn --cwd packages/web watch",
"serve": "run-p serve:*",
"docs": "cd docs && ./node_modules/.bin/astro telemetry disable && ./node_modules/.bin/astro dev --host",
Expand Down
4 changes: 4 additions & 0 deletions packages/cli/src/cli.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -389,6 +389,10 @@ export async function cli() {
"--remote-install",
"Install dependencies from remote repository"
)
.option(
"--dispatch-progress",
"Dispatch progress events to all clients"
)
.action(startServer) // Action to start the server

// Define 'parse' command group for parsing tasks
Expand Down
1 change: 0 additions & 1 deletion packages/cli/src/run.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -143,7 +143,6 @@ export async function runScriptInternal(
CancellationOptions & {
outputTrace?: MarkdownTrace
cli?: boolean
workspaceFiles?: WorkspaceFile[]
infoCb?: (partialResponse: { text: string }) => void
partialCb?: (progress: ChatCompletionsProgressReport) => void
}
Expand Down
147 changes: 86 additions & 61 deletions packages/cli/src/server.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,8 @@ import {
UNHANDLED_ERROR_CODE,
MODEL_PROVIDER_CLIENT,
} from "../../core/src/constants"
import {
isCancelError,
errorMessage,
serializeError,
} from "../../core/src/error"
import { ServerResponse, host, runtimeHost } from "../../core/src/host"
import { isCancelError, serializeError } from "../../core/src/error"
import { host, runtimeHost } from "../../core/src/host"
import { MarkdownTrace, TraceChunkEvent } from "../../core/src/trace"
import {
logVerbose,
Expand All @@ -37,8 +33,8 @@ import {
ResponseStatus,
LanguageModelConfiguration,
ServerEnvResponse,
ServerResponse,
} from "../../core/src/server/messages"
import { envInfo } from "./info"
import { LanguageModel } from "../../core/src/chat"
import {
ChatCompletionResponse,
Expand Down Expand Up @@ -70,13 +66,15 @@ export async function startServer(options: {
remoteBranch?: string
remoteForce?: boolean
remoteInstall?: boolean
dispatchProgress?: boolean
}) {
// Parse and set the server port, using a default if not specified.
const corsOrigin = options.cors || process.env.GENAISCRIPT_CORS_ORIGIN
const port = parseInt(options.port) || SERVER_PORT
const apiKey = options.apiKey || process.env.GENAISCRIPT_API_KEY
const serverHost = options.network ? "0.0.0.0" : "127.0.0.1"
const remote = options.remote
const dispatchProgress = !!options.dispatchProgress

// store original working directory
const cwd = process.cwd()
Expand Down Expand Up @@ -143,17 +141,19 @@ export async function startServer(options: {

const checkApiKey = (req: http.IncomingMessage) => {
if (!apiKey) return true
if (req.headers.authorization !== apiKey) return true

const { authorization } = req.headers
if (authorization === apiKey) return true

const url = req.url.replace(/^[^\?]*\?/, "")
const search = new URLSearchParams(url)
const hash = search.get("api-key")
if (hash === apiKey) return true

logError(`clients: connection unauthorized ${url}, ${hash}`)
logVerbose(`url:${req.url}`)
logVerbose(`api:${apiKey}`)
logVerbose(`auth:${req.headers.authorization}`)
logError(`clients: connection unauthorized ${url}`)
logVerbose(`url :${req.url}`)
logVerbose(`key :${apiKey}`)
logVerbose(`auth:${authorization}`)
logVerbose(`hash:${hash}`)
return false
}
Expand Down Expand Up @@ -266,6 +266,45 @@ export async function startServer(options: {
logVerbose(`clients: closed (${wss.clients.size} clients)`)
)

const send = (payload: object) => {
const cmsg = JSON.stringify(payload)
if (dispatchProgress)
for (const client of this.clients) client.send(cmsg)
else ws?.send(cmsg)
}
const sendProgress = (
runId: string,
payload: Omit<PromptScriptProgressResponseEvent, "type" | "runId">
) => {
send({
type: "script.progress",
runId,
...payload,
} satisfies PromptScriptProgressResponseEvent)
}

// send traces of in-flight runs
for (const [runId, run] of Object.entries(runs)) {
chunkString(run.outputTrace.content).forEach((c) =>
ws.send(
JSON.stringify({
type: "script.progress",
runId,
trace: c,
} satisfies PromptScriptProgressResponseEvent)
)
)
chunkString(run.trace.content).forEach((c) =>
ws.send(
JSON.stringify({
type: "script.progress",
runId,
trace: c,
} satisfies PromptScriptProgressResponseEvent)
)
)
}

// Handle incoming messages based on their type.
ws.on("message", async (msg) => {
const data = JSON.parse(msg.toString()) as RequestMessages
Expand Down Expand Up @@ -340,31 +379,16 @@ export async function startServer(options: {
new AbortSignalCancellationController()
const trace = new MarkdownTrace()
const outputTrace = new MarkdownTrace()
const send = (
payload: Omit<
PromptScriptProgressResponseEvent,
"type" | "runId"
>
) =>
ws?.send(
JSON.stringify(<
PromptScriptProgressResponseEvent
>{
type: "script.progress",
runId,
...payload,
})
)
trace.addEventListener(TRACE_CHUNK, (ev) => {
const tev = ev as TraceChunkEvent
chunkString(tev.chunk, 2 << 14).forEach((c) =>
send({ trace: c })
sendProgress(runId, { trace: c })
)
})
outputTrace.addEventListener(TRACE_CHUNK, (ev) => {
const tev = ev as TraceChunkEvent
chunkString(tev.chunk, 2 << 14).forEach((c) =>
send({ output: c })
sendProgress(runId, { output: c })
)
})
logVerbose(`run ${runId}: starting ${script}`)
Expand All @@ -375,15 +399,15 @@ export async function startServer(options: {
outputTrace,
cancellationToken: canceller.token,
infoCb: ({ text }) => {
send({ progress: text })
sendProgress(runId, { progress: text })
},
partialCb: ({
responseChunk,
responseSoFar,
tokensSoFar,
responseTokens,
}) => {
send({
sendProgress(runId, {
response: responseSoFar,
responseChunk,
tokens: tokensSoFar,
Expand All @@ -396,33 +420,29 @@ export async function startServer(options: {
logVerbose(
`\nrun ${runId}: completed with ${exitCode}`
)
ws?.send(
JSON.stringify(<
PromptScriptEndResponseEvent
>{
type: "script.end",
runId,
exitCode,
result,
})
)
send({
type: "script.end",
runId,
exitCode,
result,
} satisfies PromptScriptEndResponseEvent)
})
.catch((e) => {
if (canceller.controller.signal.aborted) return
if (!isCancelError(e)) trace.error(e)
logError(`\nrun ${runId}: failed`)
logError(e)
ws?.send(
JSON.stringify(<
PromptScriptEndResponseEvent
>{
type: "script.end",
runId,
exitCode: isCancelError(e)
? USER_CANCELLED_ERROR_CODE
: UNHANDLED_ERROR_CODE,
})
)
send({
type: "script.end",
runId,
result: {
status: "error",
error: serializeError(e),
},
exitCode: isCancelError(e)
? USER_CANCELLED_ERROR_CODE
: UNHANDLED_ERROR_CODE,
} satisfies PromptScriptEndResponseEvent)
})
runs[runId] = {
runner,
Expand Down Expand Up @@ -467,11 +487,21 @@ export async function startServer(options: {
} finally {
assert(!!response)
if (response.error) logError(response.error)
ws.send(JSON.stringify({ id, response }))
send({ id, type, response })
}
})
})

const setCORSHeaders = (res: http.ServerResponse) => {
res.setHeader("Access-Control-Allow-Origin", corsOrigin)
res.setHeader("Access-Control-Allow-Methods", "OPTIONS, GET")
res.setHeader("Access-Control-Max-Age", 24 * 3600) // 1 day
res.setHeader(
"Access-Control-Allow-Headers",
"Content-Type, Authorization, Accept"
)
}

// Create an HTTP server to handle basic requests.
const httpServer = http.createServer(async (req, res) => {
const { url, method } = req
Expand All @@ -482,19 +512,14 @@ export async function startServer(options: {
res.statusCode = 405
res.end()
} else {
res.setHeader("Access-Control-Allow-Origin", corsOrigin)
res.setHeader("Access-Control-Allow-Methods", "OPTIONS, GET")
res.setHeader("Access-Control-Max-Age", 24 * 3600) // 1 day
res.setHeader(
"Access-Control-Allow-Headers",
"Content-Type, Authorization, Accept"
)
setCORSHeaders(res)
res.statusCode = 204
res.end()
}
return
}

if (corsOrigin) setCORSHeaders(res)
res.setHeader("Cache-Control", "no-store")
if (method === "GET" && route === "/") {
res.setHeader("Content-Type", "text/html")
Expand Down Expand Up @@ -570,7 +595,7 @@ export async function startServer(options: {
} else socket.destroy()
})
// Start the HTTP server on the specified port.
const serverhash = apiKey ? `#api-key:${encodeURIComponent(apiKey)}` : ""
const serverhash = apiKey ? `?api-key:${encodeURIComponent(apiKey)}` : ""
httpServer.listen(port, serverHost, () => {
console.log(`GenAIScript server v${CORE_VERSION}`)
console.log(`┃ Local http://${serverHost}:${port}/${serverhash}`)
Expand Down
13 changes: 13 additions & 0 deletions packages/core/src/assert.ts
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
export function assert(
cond: boolean,
msg = "Assertion failed",
// eslint-disable-next-line @typescript-eslint/no-explicit-any
debugData?: any
) {
if (!cond) {
if (debugData) console.error(msg || `assertion failed`, debugData)
// eslint-disable-next-line no-debugger
debugger
throw new Error(msg)
}
}
4 changes: 4 additions & 0 deletions packages/core/src/constants.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,10 @@ export const TRACE_CHUNK = "traceChunk"
export const TRACE_DETAILS = "traceDetails"
export const RECONNECT = "reconnect"
export const OPEN = "open"
export const CLOSE = "close"
export const MESSAGE = "message"
export const ERROR = "error"
export const CONNECT = "connect"
export const MAX_TOOL_CALLS = 10000

// https://learn.microsoft.com/en-us/azure/ai-services/openai/reference
Expand Down
8 changes: 0 additions & 8 deletions packages/core/src/host.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,14 +51,6 @@ export interface RetrievalService {
): Promise<RetrievalSearchResponse>
}

export interface ServerResponse extends ResponseStatus {
version: string
node: string
platform: string
arch: string
pid: number
}

export interface ServerManager {
start(): Promise<void>
close(): Promise<void>
Expand Down
Loading

0 comments on commit c82caf5

Please sign in to comment.