|
| 1 | +import { createSimpleContext } from "@cyberstrike-io/ui/context" |
| 2 | +import { useParams } from "@solidjs/router" |
| 3 | +import { createEffect, createMemo, onCleanup, untrack } from "solid-js" |
| 4 | +import { createStore, produce, reconcile } from "solid-js/store" |
| 5 | +import { useSDK } from "@/context/sdk" |
| 6 | +import { |
| 7 | + activityChannels, |
| 8 | + activityRefreshChannels, |
| 9 | + isActivity, |
| 10 | + mergeActivity, |
| 11 | + type Activity, |
| 12 | + type ActivitySource, |
| 13 | + type WorkbenchChannel, |
| 14 | +} from "@/pages/session/activity" |
| 15 | + |
| 16 | +const blank = () => ({ |
| 17 | + activity: 0, |
| 18 | + mission: 0, |
| 19 | + topology: 0, |
| 20 | + memory: 0, |
| 21 | + mcp: 0, |
| 22 | + bolt: 0, |
| 23 | + terminal: 0, |
| 24 | + vulns: 0, |
| 25 | + web: 0, |
| 26 | +}) |
| 27 | + |
| 28 | +export const { use: useWorkbench, provider: WorkbenchProvider } = createSimpleContext({ |
| 29 | + name: "Workbench", |
| 30 | + init: () => { |
| 31 | + const params = useParams() |
| 32 | + const sdk = useSDK() |
| 33 | + const [events, setEvents] = createStore<Activity[]>([]) |
| 34 | + const [state, setState] = createStore({ |
| 35 | + connected: false, |
| 36 | + error: "", |
| 37 | + changes: blank(), |
| 38 | + revision: blank(), |
| 39 | + last: blank(), |
| 40 | + }) |
| 41 | + |
| 42 | + const bump = (channel: WorkbenchChannel, time = Date.now(), changed = true) => { |
| 43 | + if (changed) setState("changes", channel, (value) => value + 1) |
| 44 | + setState("revision", channel, (value) => value + 1) |
| 45 | + setState("last", channel, time) |
| 46 | + } |
| 47 | + |
| 48 | + const marked = new Set<string>() |
| 49 | + const mark = (event: Activity) => { |
| 50 | + if (marked.has(event.id)) return |
| 51 | + marked.delete(event.id) |
| 52 | + marked.add(event.id) |
| 53 | + while (marked.size > 2_000) marked.delete(marked.values().next().value!) |
| 54 | + for (const channel of activityChannels(event)) bump(channel, event.time) |
| 55 | + for (const channel of activityRefreshChannels(event)) bump(channel, event.time, false) |
| 56 | + } |
| 57 | + |
| 58 | + const add = (event: Activity) => { |
| 59 | + const index = events.findIndex((item) => item.id === event.id) |
| 60 | + if (index !== -1) { |
| 61 | + setEvents(index, reconcile(event)) |
| 62 | + mark(event) |
| 63 | + return |
| 64 | + } |
| 65 | + setEvents( |
| 66 | + produce((draft) => { |
| 67 | + draft.push(event) |
| 68 | + if (draft.length > 2_000) draft.splice(0, draft.length - 2_000) |
| 69 | + }), |
| 70 | + ) |
| 71 | + mark(event) |
| 72 | + } |
| 73 | + |
| 74 | + createEffect(() => { |
| 75 | + const sessionID = params.id |
| 76 | + sdk.directory |
| 77 | + marked.clear() |
| 78 | + setEvents(reconcile([])) |
| 79 | + setState({ |
| 80 | + connected: false, |
| 81 | + error: "", |
| 82 | + changes: blank(), |
| 83 | + revision: blank(), |
| 84 | + last: blank(), |
| 85 | + }) |
| 86 | + if (!sessionID) return |
| 87 | + |
| 88 | + const abort = new AbortController() |
| 89 | + const client = sdk.createClient({ |
| 90 | + directory: sdk.directory, |
| 91 | + throwOnError: true, |
| 92 | + signal: abort.signal, |
| 93 | + }) |
| 94 | + let connected = false |
| 95 | + let hydrated = false |
| 96 | + let requested = 0 |
| 97 | + let completed = 0 |
| 98 | + let syncing: Promise<void> | undefined |
| 99 | + let streamError = "" |
| 100 | + let historyError = "" |
| 101 | + const render = () => setState({ connected, error: historyError || streamError }) |
| 102 | + const snapshot = async () => { |
| 103 | + const known = new Set(untrack(() => events.map((event) => event.id))) |
| 104 | + const pages: Activity[] = [] |
| 105 | + let before: number | undefined |
| 106 | + let beforeID: string | undefined |
| 107 | + while (pages.length < 2_000) { |
| 108 | + const limit = Math.min(500, 2_000 - pages.length) |
| 109 | + const response = await client.eventLog.list({ sessionID, before, beforeID, limit }) |
| 110 | + const page = (response.data ?? []).filter(isActivity) |
| 111 | + if (page.length === 0) break |
| 112 | + pages.unshift(...page) |
| 113 | + if (hydrated && page.some((event) => known.has(event.id))) break |
| 114 | + if (page.length < limit) break |
| 115 | + const next = page[0]!.time |
| 116 | + const nextID = page[0]!.id |
| 117 | + if (before === next && beforeID === nextID) break |
| 118 | + before = next |
| 119 | + beforeID = nextID |
| 120 | + } |
| 121 | + return { known, incoming: mergeActivity(pages, [], 2_000) } |
| 122 | + } |
| 123 | + const sync = () => { |
| 124 | + if (syncing) return syncing |
| 125 | + syncing = (async () => { |
| 126 | + while (!abort.signal.aborted) { |
| 127 | + const version = requested |
| 128 | + try { |
| 129 | + const history = await snapshot() |
| 130 | + if (abort.signal.aborted) return |
| 131 | + const changed = hydrated |
| 132 | + setEvents(reconcile(mergeActivity(history.incoming, [...events]))) |
| 133 | + if (changed) history.incoming.filter((event) => !history.known.has(event.id)).forEach(mark) |
| 134 | + hydrated = true |
| 135 | + completed = version |
| 136 | + historyError = "" |
| 137 | + render() |
| 138 | + if (requested <= completed) return |
| 139 | + } catch (cause) { |
| 140 | + if (abort.signal.aborted) return |
| 141 | + historyError = cause instanceof Error ? cause.message : String(cause) |
| 142 | + render() |
| 143 | + await new Promise((resolve) => setTimeout(resolve, 500)) |
| 144 | + } |
| 145 | + } |
| 146 | + })().finally(() => { |
| 147 | + syncing = undefined |
| 148 | + }) |
| 149 | + return syncing |
| 150 | + } |
| 151 | + const requestSync = () => { |
| 152 | + requested++ |
| 153 | + return sync() |
| 154 | + } |
| 155 | + const unsubs = [ |
| 156 | + sdk.event.on("memory.updated", (event) => { |
| 157 | + if (!event.properties.sessionID) bump("memory") |
| 158 | + }), |
| 159 | + sdk.event.on("mcp.tools.changed", () => bump("mcp")), |
| 160 | + sdk.event.on("pty.created", () => bump("terminal")), |
| 161 | + sdk.event.on("pty.updated", () => bump("terminal")), |
| 162 | + sdk.event.on("pty.exited", () => bump("terminal")), |
| 163 | + sdk.event.on("pty.deleted", () => bump("terminal")), |
| 164 | + ] |
| 165 | + |
| 166 | + void requestSync() |
| 167 | + let timer: ReturnType<typeof setTimeout> | undefined |
| 168 | + void (async () => { |
| 169 | + while (!abort.signal.aborted) { |
| 170 | + try { |
| 171 | + const response = await client.eventLog.stream( |
| 172 | + { sessionID }, |
| 173 | + { |
| 174 | + onSseError: (cause) => { |
| 175 | + if (abort.signal.aborted) return |
| 176 | + connected = false |
| 177 | + streamError = cause instanceof Error ? cause.message : String(cause) |
| 178 | + render() |
| 179 | + }, |
| 180 | + onSseEvent: () => { |
| 181 | + if (timer) clearTimeout(timer) |
| 182 | + timer = undefined |
| 183 | + const recovered = !connected |
| 184 | + connected = true |
| 185 | + streamError = "" |
| 186 | + render() |
| 187 | + if (recovered) void requestSync() |
| 188 | + }, |
| 189 | + }, |
| 190 | + ) |
| 191 | + timer = setTimeout(() => { |
| 192 | + if (abort.signal.aborted) return |
| 193 | + streamError = "Live activity stream did not connect" |
| 194 | + render() |
| 195 | + }, 5_000) |
| 196 | + let next = await response.stream.next() |
| 197 | + while (!next.done && !abort.signal.aborted) { |
| 198 | + if (isActivity(next.value)) add(next.value) |
| 199 | + next = await response.stream.next() |
| 200 | + } |
| 201 | + if (abort.signal.aborted) return |
| 202 | + connected = false |
| 203 | + streamError = "Live activity stream disconnected" |
| 204 | + render() |
| 205 | + } catch (cause) { |
| 206 | + if (abort.signal.aborted) return |
| 207 | + connected = false |
| 208 | + streamError = cause instanceof Error ? cause.message : String(cause) |
| 209 | + render() |
| 210 | + } finally { |
| 211 | + if (timer) clearTimeout(timer) |
| 212 | + } |
| 213 | + await new Promise((resolve) => setTimeout(resolve, 500)) |
| 214 | + } |
| 215 | + })() |
| 216 | + |
| 217 | + onCleanup(() => { |
| 218 | + if (timer) clearTimeout(timer) |
| 219 | + abort.abort() |
| 220 | + unsubs.forEach((unsub) => unsub()) |
| 221 | + }) |
| 222 | + }) |
| 223 | + |
| 224 | + const latest = createMemo(() => events.at(-1)) |
| 225 | + |
| 226 | + return { |
| 227 | + get events() { |
| 228 | + return events |
| 229 | + }, |
| 230 | + get connected() { |
| 231 | + return state.connected |
| 232 | + }, |
| 233 | + get error() { |
| 234 | + return state.error |
| 235 | + }, |
| 236 | + latest, |
| 237 | + count(source: ActivitySource) { |
| 238 | + return events.filter((event) => event.source === source).length |
| 239 | + }, |
| 240 | + changes(channel: WorkbenchChannel) { |
| 241 | + return state.changes[channel] |
| 242 | + }, |
| 243 | + revision(channel: WorkbenchChannel) { |
| 244 | + return state.revision[channel] |
| 245 | + }, |
| 246 | + last(channel: WorkbenchChannel) { |
| 247 | + return state.last[channel] |
| 248 | + }, |
| 249 | + ack(channel: WorkbenchChannel) { |
| 250 | + setState("changes", channel, 0) |
| 251 | + }, |
| 252 | + } |
| 253 | + }, |
| 254 | +}) |
0 commit comments