diff --git a/apps/admin-web/src/components/PublishArticleModal.vue b/apps/admin-web/src/components/PublishArticleModal.vue index 5f180d1..b8c7b54 100644 --- a/apps/admin-web/src/components/PublishArticleModal.vue +++ b/apps/admin-web/src/components/PublishArticleModal.vue @@ -11,7 +11,12 @@ import { articlesApi, mediaApi, publishJobsApi, resolveApiURL, tenantAccountsApi import { coverUploadRequired, deriveCoverFileName } from "@/lib/cover-requirements"; import { formatDateTime } from "@/lib/display"; import { formatError } from "@/lib/errors"; -import { getPublishPlatformMeta, normalizePublishPlatformId, normalizePublishPlatformIds } from "@/lib/publish-platforms"; +import { + getPublishPlatformMeta, + isPublishPlatformId, + normalizePublishPlatformId, + normalizePublishPlatformIds, +} from "@/lib/publish-platforms"; type PublishState = "immediate" | "queued" | "unavailable"; @@ -109,8 +114,19 @@ const platformMap = computed(() => { ); }); +const publishAccounts = computed(() => { + return (accountsQuery.data.value ?? []).filter((account) => { + if (account.deleted_at) { + return false; + } + + const platformId = normalizePublishPlatformId(account.platform); + return platformMap.value.has(platformId) || isPublishPlatformId(platformId); + }); +}); + const rawAccountMap = computed(() => { - return new Map((accountsQuery.data.value ?? []).map((account) => [account.id, account])); + return new Map(publishAccounts.value.map((account) => [account.id, account])); }); const articlePlatformIds = computed(() => normalizePublishPlatformIds(detailQuery.data.value?.platforms ?? [])); @@ -118,8 +134,7 @@ const articlePlatformIds = computed(() => normalizePublishPlatformIds(detailQuer const accountCards = computed(() => { const preferredPlatforms = new Set(articlePlatformIds.value); - return [...(accountsQuery.data.value ?? [])] - .filter((account) => !account.deleted_at) + return [...publishAccounts.value] .map((account) => { const platformId = normalizePublishPlatformId(account.platform); const platform = platformMap.value.get(platformId); diff --git a/apps/admin-web/src/i18n/messages/zh-CN.ts b/apps/admin-web/src/i18n/messages/zh-CN.ts index 0d0e57e..8b11f9b 100644 --- a/apps/admin-web/src/i18n/messages/zh-CN.ts +++ b/apps/admin-web/src/i18n/messages/zh-CN.ts @@ -349,7 +349,7 @@ const zhCN = { collectBrandRequired: "请先选择品牌，再发起采集", collectKeywordRequired: "请先选择关键词，再立即采集当前问题", collectClientChecking: "正在检查当前账号的桌面客户端状态", - collectClientOnlineRequired: "当前登录账号的 desktop-client 未在线，暂时不能立即采集", + collectClientOnlineRequired: "当前登录账号的桌面客户端未在线，暂时不能立即采集", pluginMode: "插件采样", pluginHint: "插件模式下结果受设备在线率影响，低覆盖样本仅供趋势参考。", awaitingSnapshot: "等待下一次有效快照", diff --git a/apps/admin-web/src/lib/errors.ts b/apps/admin-web/src/lib/errors.ts index 5376ede..06c861c 100644 --- a/apps/admin-web/src/lib/errors.ts +++ b/apps/admin-web/src/lib/errors.ts @@ -59,12 +59,7 @@ const errorMessageMap: Record = { publisher_plugin_empty_response: "浏览器插件未返回数据，请刷新当前页面后重试", publisher_plugin_invalid_response: "浏览器插件返回了无效数据，请刷新当前页面后重试", publisher_plugin_unknown_error: "浏览器插件调用失败，请稍后重试", - monitoring_plugin_timeout: "监测插件响应超时，请稍后重试", - monitoring_plugin_invalid_response: "监测插件返回了无效数据，请稍后重试", - monitoring_plugin_unknown_error: "监测插件调用失败，请稍后重试", - monitoring_plugin_required: "请先安装并打开浏览器插件", - monitoring_plugin_register_required: "浏览器插件正在初始化，请稍后重试", - installation_offline: "监测插件离线，请打开浏览器插件后重试", + installation_offline: "桌面客户端未在线，请打开桌面客户端后重试", network_error: "网络连接失败，请稍后重试", unknown_error: "发生未知错误", }; diff --git a/apps/admin-web/src/lib/monitoring-plugin.ts b/apps/admin-web/src/lib/monitoring-plugin.ts deleted file mode 100644 index 479fb94..0000000 --- a/apps/admin-web/src/lib/monitoring-plugin.ts +++ /dev/null @@ -1,155 +0,0 @@ -import type { - MonitoringCycleKickoffResponse, - MonitoringCycleResult, - MonitoringHeartbeatSyncResult, - MonitoringLocalPlatformState, - MonitoringPluginPingResponse, - MonitoringRuntimeStateSnapshot, -} from "@geo/shared-types"; - -type MonitoringAction = "ping" | "detectPlatforms" | "syncHeartbeat" | "runCycle" | "kickoffCycle" | "getState"; - -type MonitoringPluginSuccessPayload = { - source: "geo-extension"; - type: "geo.monitoring.response"; - requestId: string; - success: true; - data: T; -}; - -type MonitoringPluginErrorPayload = { - source: "geo-extension"; - type: "geo.monitoring.response"; - requestId: string; - success: false; - error: string; -}; - -type MonitoringPluginPayloadMap = { - ping: MonitoringPluginPingResponse; - detectPlatforms: MonitoringLocalPlatformState[]; - syncHeartbeat: MonitoringHeartbeatSyncResult; - runCycle: MonitoringCycleResult; - kickoffCycle: MonitoringCycleKickoffResponse; - getState: MonitoringRuntimeStateSnapshot; -}; - -function isPlainObject(value: unknown): value is Record { - return typeof value === "object" && value !== null; -} - -function isMonitoringPingPayload(value: unknown): value is MonitoringPluginPingResponse { - return isPlainObject(value) && typeof value.installed === "boolean"; -} - -function isMonitoringDetectPayload(value: unknown): value is MonitoringLocalPlatformState[] { - return Array.isArray(value); -} - -function isMonitoringHeartbeatPayload(value: unknown): value is MonitoringHeartbeatSyncResult { - return isPlainObject(value) && typeof value.success === "boolean" && Array.isArray(value.platform_states); -} - -function isMonitoringCyclePayload(value: unknown): value is MonitoringCycleResult { - return ( - isPlainObject(value) && - typeof value.success === "boolean" && - typeof value.leased_task_count === "number" && - typeof value.processed_task_count === "number" && - typeof value.failed_task_count === "number" && - typeof value.skipped_task_count === "number" - ); -} - -function isMonitoringKickoffPayload(value: unknown): value is MonitoringCycleKickoffResponse { - return isPlainObject(value) && value.accepted === true && typeof value.status === "string"; -} - -function isMonitoringStatePayload(value: unknown): value is MonitoringRuntimeStateSnapshot { - return isPlainObject(value) && typeof value.installation_key === "string" && Array.isArray(value.monitoring_platforms); -} - -function isMonitoringPayloadValid(action: T, value: unknown): value is MonitoringPluginPayloadMap[T] { - switch (action) { - case "ping": - return isMonitoringPingPayload(value); - case "detectPlatforms": - return isMonitoringDetectPayload(value); - case "syncHeartbeat": - return isMonitoringHeartbeatPayload(value); - case "runCycle": - return isMonitoringCyclePayload(value); - case "kickoffCycle": - return isMonitoringKickoffPayload(value); - case "getState": - return isMonitoringStatePayload(value); - default: - return false; - } -} - -function nextRequestId(): string { - if (typeof crypto !== "undefined" && "randomUUID" in crypto) { - return crypto.randomUUID(); - } - return `geo-monitoring-${Date.now()}-${Math.random().toString(36).slice(2)}`; -} - -async function requestMonitoringPlugin(action: T, timeoutMs = 5000): Promise { - if (typeof window === "undefined") { - throw new Error("plugin bridge is only available in browser contexts"); - } - - const requestId = nextRequestId(); - - return new Promise((resolve, reject) => { - const timer = window.setTimeout(() => { - cleanup(); - reject(new Error("monitoring_plugin_timeout")); - }, timeoutMs); - - const handleMessage = (event: MessageEvent | MonitoringPluginErrorPayload>) => { - if (event.source !== window || !event.data || event.data.type !== "geo.monitoring.response") { - return; - } - if (event.data.requestId !== requestId || event.data.source !== "geo-extension") { - return; - } - - cleanup(); - if (event.data.success) { - if (!isMonitoringPayloadValid(action, event.data.data)) { - reject(new Error("monitoring_plugin_invalid_response")); - return; - } - resolve(event.data.data); - return; - } - reject(new Error(event.data.error || "monitoring_plugin_unknown_error")); - }; - - const cleanup = () => { - window.clearTimeout(timer); - window.removeEventListener("message", handleMessage); - }; - - window.addEventListener("message", handleMessage); - window.postMessage( - { - source: "geo-admin", - type: "geo.monitoring.request", - requestId, - action, - }, - window.location.origin, - ); - }); -} - -export async function runMonitoringCycle(): Promise { - return requestMonitoringPlugin("runCycle", 30000); -} - -export async function kickoffMonitoringCycle(): Promise { - return requestMonitoringPlugin("kickoffCycle", 5000); -} diff --git a/apps/desktop-client/package.json b/apps/desktop-client/package.json index 0cb2056..808955a 100644 --- a/apps/desktop-client/package.json +++ b/apps/desktop-client/package.json @@ -24,11 +24,13 @@ "pino": "^9.0.0", "playwright-core": "^1.55.0", "vue": "^3.5.31", - "vue-router": "^4.5.1" + "vue-router": "^4.5.1", + "ws": "^8.20.0" }, "devDependencies": { "@playwright/test": "^1.0.0", "@types/node": "^24.0.0", + "@types/ws": "^8.18.1", "@vitejs/plugin-vue": "^5.2.4", "electron": "41.2.0", "electron-builder": "^25.0.0", diff --git a/apps/desktop-client/src/main/runtime-controller.ts b/apps/desktop-client/src/main/runtime-controller.ts index 0987520..9fac1a0 100644 --- a/apps/desktop-client/src/main/runtime-controller.ts +++ b/apps/desktop-client/src/main/runtime-controller.ts @@ -89,6 +89,7 @@ import { upsertDesktopAccount, } from "./transport/api-client"; import { SseClient, SseClientError } from "./transport/sse-client"; +import { DispatchWsClient } from "./transport/dispatch-ws-client"; type RuntimeTone = "info" | "success" | "warn" | "danger"; type RuntimeTaskStatus = @@ -184,6 +185,8 @@ interface RuntimeState { heartbeatTimer: ReturnType | null; pullTimer: ReturnType | null; sseClient: SseClient | null; + dispatchWsClient: DispatchWsClient | null; + dispatchWsConnected: boolean; lastHeartbeatAt: number; lastHeartbeatStatus: "idle" | "success" | "failed"; lastPullAt: number; @@ -210,6 +213,8 @@ const state: RuntimeState = { heartbeatTimer: null, pullTimer: null, sseClient: null, + dispatchWsClient: null, + dispatchWsConnected: false, lastHeartbeatAt: 0, lastHeartbeatStatus: "idle", lastPullAt: 0, @@ -354,6 +359,7 @@ function startRuntime(): void { recordActivity("success", "调度节点已启动", "已启动 SSE、心跳上报和兜底拉取。"); connectEventStream(); + connectDispatchWs(); scheduleHeartbeatLoop(); schedulePullLoop(); @@ -383,6 +389,11 @@ function stopRuntime(): void { state.sseClient.stop(); state.sseClient = null; } + if (state.dispatchWsClient) { + state.dispatchWsClient.stop(); + state.dispatchWsClient = null; + } + state.dispatchWsConnected = false; setTransportSseState("idle"); setSchedulerPhase("idle"); @@ -521,6 +532,63 @@ function connectEventStream(): void { sseClient.start(); } +function connectDispatchWs(): void { + if (!canRunLive(state.session)) { + return; + } + + state.dispatchWsClient?.stop(); + + const wsUrl = buildApiUrl(state.session.api_base_url, "/api/desktop/dispatch") + .replace(/^http:\/\//, "ws://") + .replace(/^https:\/\//, "wss://"); + + const client = new DispatchWsClient({ + url: wsUrl, + headers: { + Authorization: `Bearer ${state.session.client_token}`, + }, + }); + + client.on("open", () => { + state.dispatchWsConnected = true; + setSchedulerError(null); + recordActivity("success", "任务分发通道已连接", "AMQP 直推 WebSocket 已就绪,发布任务将以低延迟到达。"); + }); + + client.on("connected", () => { + state.dispatchWsConnected = true; + }); + + client.on("task_available", (payload) => { + if (!isDesktopTaskEvent(payload)) { + return; + } + state.dispatchWsConnected = true; + handleTaskEvent(payload); + }); + + client.on("error", (payload) => { + state.dispatchWsConnected = false; + const message = errorMessage(payload); + recordActivity("warn", "任务分发通道异常", message); + }); + + client.on("close", () => { + state.dispatchWsConnected = false; + }); + + client.on("reconnect", (payload) => { + if (!isReconnectPayload(payload)) { + return; + } + recordActivity("info", "任务分发通道退避等待", `将在 ${payload.delayMs}ms 后重连 WebSocket。`); + }); + + state.dispatchWsClient = client; + client.start(); +} + function handleTaskEvent(event: DesktopTaskEventMessage): void { const existing = state.tasks.get(event.task_id); const next: RuntimeTaskRecord = { @@ -871,11 +939,18 @@ function pumpExecutionLoop(): void { } if (state.activeExecutions.size === 0) { - void pullNextTask("publish"); + // Publish task_available is push-delivered over the dispatch WebSocket. We + // only fall back to the HTTP pull when the push channel is down, otherwise + // every scheduler tick would wake the /lease endpoint for no reason. + if (!state.dispatchWsConnected) { + void pullNextTask("publish"); + } return; } if (canStartAnotherMonitorExecution()) { + // Monitor has not been migrated to the dispatch channel yet, so it still + // relies on the periodic pull as its primary signal. void pullNextTask("monitor"); } } diff --git a/apps/desktop-client/src/main/transport/dispatch-ws-client.ts b/apps/desktop-client/src/main/transport/dispatch-ws-client.ts new file mode 100644 index 0000000..8567bbd --- /dev/null +++ b/apps/desktop-client/src/main/transport/dispatch-ws-client.ts @@ -0,0 +1,211 @@ +import WebSocket, { type RawData } from "ws"; + +export type DispatchEventHandler = (payload: unknown) => void; + +export interface DispatchWsClientOptions { + url: string; + headers?: Record; + initialRetryDelayMs?: number; + maxRetryDelayMs?: number; +} + +type DispatchWsState = "idle" | "connecting" | "open"; + +/** + * DispatchWsClient receives per-client task dispatch messages from the server + * over a single authenticated WebSocket. It mirrors {@link SseClient} so the + * runtime-controller can treat it as a drop-in replacement for the publish + * side of the SSE channel. + * + * The underlying contract is intentionally narrow: + * - The server pushes newline-free JSON frames whose shape matches + * `DesktopDispatchEvent` on the Go side. + * - The client never sends application frames back; acknowledgement is done + * out-of-band via `/tasks/lease`. + * + * Reconnect uses exponential backoff with jitter. On every fresh TCP + * connection the server replays a synthetic "connected" frame so callers can + * observe the auth + transport handshake is healthy. + */ +export class DispatchWsClient { + readonly #listeners = new Map>(); + readonly #options: DispatchWsClientOptions; + #running = false; + #socket: WebSocket | null = null; + #retryDelayMs: number; + #reconnectTimer: ReturnType | null = null; + + constructor(options: DispatchWsClientOptions) { + this.#options = options; + this.#retryDelayMs = options.initialRetryDelayMs ?? 1_000; + } + + on(event: string, handler: DispatchEventHandler): () => void { + const bucket = this.#listeners.get(event) ?? new Set(); + bucket.add(handler); + this.#listeners.set(event, bucket); + return () => { + bucket.delete(handler); + if (bucket.size === 0) { + this.#listeners.delete(event); + } + }; + } + + start(): void { + if (this.#running) { + return; + } + this.#running = true; + this.#connect(); + } + + stop(): void { + this.#running = false; + this.#clearReconnect(); + this.#setState("idle"); + + const socket = this.#socket; + this.#socket = null; + if (!socket) { + return; + } + + // Detach every listener we registered so their handlers don't re-fire our + // own reconnect / error callbacks during teardown, but keep a noop error + // listener around: `ws` emits an asynchronous "error" when a socket still + // in the CONNECTING state is closed, and without a listener Node turns it + // into an uncaughtException. + socket.removeAllListeners("open"); + socket.removeAllListeners("message"); + socket.removeAllListeners("close"); + socket.removeAllListeners("error"); + socket.on("error", () => { + // swallow: teardown path, we no longer care about the socket's fate. + }); + + try { + if (socket.readyState === WebSocket.CONNECTING) { + // `close()` on a CONNECTING socket throws asynchronously; `terminate()` + // destroys the underlying TCP connection and is always safe. + socket.terminate(); + } else if (socket.readyState === WebSocket.OPEN) { + socket.close(1000, "client_stop"); + } + } catch { + // Fallback: ignore; the reference is already detached. + } + } + + emit(event: string, payload: unknown): void { + for (const handler of this.#listeners.get(event) ?? []) { + try { + handler(payload); + } catch { + // Swallow handler errors; listeners must not kill the transport. + } + } + } + + #connect(): void { + if (!this.#running) { + return; + } + this.#setState("connecting"); + + let socket: WebSocket; + try { + socket = new WebSocket(this.#options.url, { + headers: this.#options.headers, + handshakeTimeout: 15_000, + }); + } catch (error) { + this.emit("error", error); + this.#scheduleReconnect(); + return; + } + + this.#socket = socket; + + socket.on("open", () => { + this.#retryDelayMs = this.#options.initialRetryDelayMs ?? 1_000; + this.#setState("open"); + this.emit("open", null); + }); + + socket.on("message", (data: RawData, isBinary: boolean) => { + if (isBinary) { + return; + } + const raw = typeof data === "string" ? data : data.toString("utf-8"); + if (!raw) { + return; + } + let parsed: unknown; + try { + parsed = JSON.parse(raw); + } catch { + this.emit("message", { event: "message", data: raw }); + return; + } + const event = extractEventType(parsed) ?? "message"; + this.emit(event, parsed); + this.emit("message", { event, data: parsed }); + }); + + socket.on("error", (error: Error) => { + this.emit("error", error); + }); + + socket.on("close", (code: number, reason: Buffer) => { + this.#socket = null; + this.emit("close", { code, reason: reason.toString() }); + if (!this.#running) { + return; + } + this.#scheduleReconnect(); + }); + } + + #scheduleReconnect(): void { + if (!this.#running) { + return; + } + this.#clearReconnect(); + + const baseDelay = this.#retryDelayMs; + const maxDelay = this.#options.maxRetryDelayMs ?? 30_000; + const jitter = Math.floor(baseDelay * 0.2 * Math.random()); + const delay = Math.min(baseDelay + jitter, maxDelay); + + this.emit("reconnect", { delayMs: delay }); + this.#reconnectTimer = globalThis.setTimeout(() => { + this.#reconnectTimer = null; + this.#connect(); + }, delay); + + this.#retryDelayMs = Math.min(this.#retryDelayMs * 2, maxDelay); + } + + #clearReconnect(): void { + if (this.#reconnectTimer) { + globalThis.clearTimeout(this.#reconnectTimer); + this.#reconnectTimer = null; + } + } + + #setState(state: DispatchWsState): void { + this.emit("state", state); + } +} + +function extractEventType(payload: unknown): string | null { + if (!payload || typeof payload !== "object") { + return null; + } + const maybe = (payload as { type?: unknown }).type; + if (typeof maybe === "string" && maybe.length > 0) { + return maybe; + } + return null; +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index fb51636..cfc0183 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -165,6 +165,9 @@ importers: vue-router: specifier: ^4.5.1 version: 4.6.4(vue@3.5.31(typescript@5.9.3)) + ws: + specifier: ^8.20.0 + version: 8.20.0 devDependencies: '@playwright/test': specifier: ^1.0.0 @@ -172,6 +175,9 @@ importers: '@types/node': specifier: ^24.0.0 version: 24.12.0 + '@types/ws': + specifier: ^8.18.1 + version: 8.18.1 '@vitejs/plugin-vue': specifier: ^5.2.4 version: 5.2.4(vite@5.4.21(@types/node@24.12.0)(lightningcss@1.32.0))(vue@3.5.31(typescript@5.9.3)) @@ -1366,6 +1372,9 @@ packages: '@types/web-bluetooth@0.0.21': resolution: {integrity: sha512-oIQLCGWtcFZy2JW77j9k8nHzAOpqMHLQejDA48XXMWH6tjCQHz5RCFz1bzsmROyL6PUm+LLnUiI4BCn221inxA==} + '@types/ws@8.18.1': + resolution: {integrity: sha512-ThVF6DCVhA8kUGy+aazFQ4kXQ7E1Ty7A3ypFOe0IcJV8O/M511G99AW24irKrW56Wt44yG9+ij8FaqoBGkuBXg==} + '@types/yauzl@2.10.3': resolution: {integrity: sha512-oJoftv0LSuaDZE3Le4DbKX+KS9G36NzOeSap90UIK0yMA/NhKJhqlSGtNDORNRaIbQfzjXDrQa0ytJ6mNRGz/Q==} @@ -4462,6 +4471,18 @@ packages: wrappy@1.0.2: resolution: {integrity: sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==} + ws@8.20.0: + resolution: {integrity: sha512-sAt8BhgNbzCtgGbt2OxmpuryO63ZoDk/sqaB/znQm94T4fCEsy/yV+7CdC1kJhOU9lboAEU7R3kquuycDoibVA==} + engines: {node: '>=10.0.0'} + peerDependencies: + bufferutil: ^4.0.1 + utf-8-validate: '>=5.0.2' + peerDependenciesMeta: + bufferutil: + optional: true + utf-8-validate: + optional: true + wsl-utils@0.3.1: resolution: {integrity: sha512-g/eziiSUNBSsdDJtCLB8bdYEUMj4jR7AGeUo96p/3dTafgjHhpF4RiCFPiRILwjQoDXx5MqkBr4fwWtR3Ky4Wg==} engines: {node: '>=20'} @@ -5923,6 +5944,10 @@ snapshots: '@types/web-bluetooth@0.0.21': {} + '@types/ws@8.18.1': + dependencies: + '@types/node': 24.12.0 + '@types/yauzl@2.10.3': dependencies: '@types/node': 24.12.0 @@ -9555,6 +9580,8 @@ snapshots: wrappy@1.0.2: {} + ws@8.20.0: {} + wsl-utils@0.3.1: dependencies: is-wsl: 3.1.1 diff --git a/server/cmd/tenant-api/main.go b/server/cmd/tenant-api/main.go index 97c7aee..57f6bac 100644 --- a/server/cmd/tenant-api/main.go +++ b/server/cmd/tenant-api/main.go @@ -30,6 +30,7 @@ func main() { app.GenerationStreams.Run(workerCtx) app.DesktopTaskStreams.Run(workerCtx) + app.DesktopDispatch.Run(workerCtx) monitoringCallbackService := tenantapp.NewMonitoringCallbackService(app.DB, app.MonitoringDB, app.RabbitMQ, app.LLM, app.Logger) tenantapp.NewMonitoringResultIngestWorker(monitoringCallbackService, app.RabbitMQ, app.Logger, app.Config.MonitoringWorkers).Start(workerCtx) diff --git a/server/configs/config.yaml b/server/configs/config.yaml index 300abcd..0204150 100644 --- a/server/configs/config.yaml +++ b/server/configs/config.yaml @@ -43,6 +43,10 @@ rabbitmq: generation_dlq: generation.task.run.dlq generation_dlq_routing_key: generation.task.run.dlq generation_event_exchange: generation.event + desktop_task_event_exchange: desktop.task.event + desktop_dispatch_exchange: desktop.task.dispatch + desktop_dispatch_publish_prefix: publish + desktop_dispatch_monitor_prefix: monitor template_assist_exchange: template.assist template_assist_routing_key: template.assist.run template_assist_queue: template.assist.run diff --git a/server/go.mod b/server/go.mod index ce6f845..9002d3d 100644 --- a/server/go.mod +++ b/server/go.mod @@ -1,6 +1,6 @@ module github.com/geo-platform/tenant-api -go 1.22.2 +go 1.23 require ( github.com/alicebob/miniredis/v2 v2.37.0 @@ -10,6 +10,7 @@ require ( github.com/gin-gonic/gin v1.9.1 github.com/golang-jwt/jwt/v5 v5.2.1 github.com/google/uuid v1.6.0 + github.com/gorilla/websocket v1.5.3 github.com/jackc/pgx/v5 v5.5.5 github.com/ledongthuc/pdf v0.0.0-20240314090751-a2a84ec735c3 github.com/minio/minio-go/v7 v7.0.83 diff --git a/server/go.sum b/server/go.sum index 3ec8e2d..c283b60 100644 --- a/server/go.sum +++ b/server/go.sum @@ -83,6 +83,8 @@ github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/ github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= +github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4= github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= @@ -223,8 +225,6 @@ golang.org/x/exp v0.0.0-20230905200255-921286631fa9 h1:GoHiUyI/Tp2nVkLI2mCxVkOjs golang.org/x/exp v0.0.0-20230905200255-921286631fa9/go.mod h1:S2oDrQGGwySpoQPVqRShND87VCbxmc6bL1Yd2oYrm6k= golang.org/x/image v0.14.0 h1:tNgSxAFe3jC4uYqvZdTr84SZoM1KfwdC9SKIFrLjFn4= golang.org/x/image v0.14.0/go.mod h1:HUYqC05R2ZcZ3ejNQsIHQDQiwWM4JBqmm6MKANTp4LE= -golang.org/x/image v0.39.0 h1:skVYidAEVKgn8lZ602XO75asgXBgLj9G/FE3RbuPFww= -golang.org/x/image v0.39.0/go.mod h1:sIbmppfU+xFLPIG0FoVUTvyBMmgng1/XAMhQ2ft0hpA= golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU= golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= @@ -240,8 +240,6 @@ golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.10.0 h1:3NQrjDixjgGwUOCaF8w2+VYHv0Ve/vGYSbdkTa98gmQ= golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= -golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= -golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= @@ -250,8 +248,6 @@ golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo= golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ= -golang.org/x/text v0.36.0 h1:JfKh3XmcRPqZPKevfXVpI1wXPTqbkE5f7JA92a55Yxg= -golang.org/x/text v0.36.0/go.mod h1:NIdBknypM8iqVmPiuco0Dh6P5Jcdk8lJL0CUebqK164= golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk= golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= diff --git a/server/internal/bootstrap/bootstrap.go b/server/internal/bootstrap/bootstrap.go index cca37af..4e1560d 100644 --- a/server/internal/bootstrap/bootstrap.go +++ b/server/internal/bootstrap/bootstrap.go @@ -45,8 +45,9 @@ type App struct { VectorStore retrieval.VectorStore ObjectStorage objectstorage.Client AuditLogs *auditlog.AsyncWriter - GenerationStreams *stream.GenerationHub + GenerationStreams *stream.GenerationHub DesktopTaskStreams *stream.DesktopTaskHub + DesktopDispatch *stream.DesktopDispatchHub Cache cache.Cache KolProfiles repository.KolProfileRepository KolPackages repository.KolPackageRepository @@ -111,6 +112,10 @@ func New(configPath string) (*App, error) { auditLogs := auditlog.NewAsyncWriter(pool, logger) generationStreams := stream.NewGenerationHub(mqClient) desktopTaskStreams := stream.NewDesktopTaskHub(mqClient) + desktopDispatchBindings := []string{ + cfg.RabbitMQ.DesktopDispatchPublishPrefix + ".*", + } + desktopDispatch := stream.NewDesktopDispatchHub(mqClient, logger, desktopDispatchBindings) appCache := cache.New(cfg.Cache.Driver, rdb) kolProfiles := repository.NewKolProfileRepository(pool) kolPackages := repository.NewKolPackageRepository(pool) @@ -170,8 +175,9 @@ func New(configPath string) (*App, error) { VectorStore: vectorStore, ObjectStorage: objectStorageClient, AuditLogs: auditLogs, - GenerationStreams: generationStreams, + GenerationStreams: generationStreams, DesktopTaskStreams: desktopTaskStreams, + DesktopDispatch: desktopDispatch, Cache: appCache, KolProfiles: kolProfiles, KolPackages: kolPackages, diff --git a/server/internal/shared/config/config.go b/server/internal/shared/config/config.go index eb03a05..00fedf9 100644 --- a/server/internal/shared/config/config.go +++ b/server/internal/shared/config/config.go @@ -78,6 +78,9 @@ type RabbitMQConfig struct { GenerationDLQRouteKey string `mapstructure:"generation_dlq_routing_key"` GenerationEventExchange string `mapstructure:"generation_event_exchange"` DesktopTaskEventExchange string `mapstructure:"desktop_task_event_exchange"` + DesktopDispatchExchange string `mapstructure:"desktop_dispatch_exchange"` + DesktopDispatchPublishPrefix string `mapstructure:"desktop_dispatch_publish_prefix"` + DesktopDispatchMonitorPrefix string `mapstructure:"desktop_dispatch_monitor_prefix"` TemplateAssistExchange string `mapstructure:"template_assist_exchange"` TemplateAssistRoutingKey string `mapstructure:"template_assist_routing_key"` TemplateAssistQueue string `mapstructure:"template_assist_queue"` @@ -475,6 +478,15 @@ func normalizeRabbitMQConfig(cfg *RabbitMQConfig) { if strings.TrimSpace(cfg.DesktopTaskEventExchange) == "" { cfg.DesktopTaskEventExchange = "desktop.task.event" } + if strings.TrimSpace(cfg.DesktopDispatchExchange) == "" { + cfg.DesktopDispatchExchange = "desktop.task.dispatch" + } + if strings.TrimSpace(cfg.DesktopDispatchPublishPrefix) == "" { + cfg.DesktopDispatchPublishPrefix = "publish" + } + if strings.TrimSpace(cfg.DesktopDispatchMonitorPrefix) == "" { + cfg.DesktopDispatchMonitorPrefix = "monitor" + } if strings.TrimSpace(cfg.TemplateAssistExchange) == "" { cfg.TemplateAssistExchange = "template.assist" } diff --git a/server/internal/shared/messaging/rabbitmq/client.go b/server/internal/shared/messaging/rabbitmq/client.go index 54d598f..36484fb 100644 --- a/server/internal/shared/messaging/rabbitmq/client.go +++ b/server/internal/shared/messaging/rabbitmq/client.go @@ -46,6 +46,16 @@ func New(_ context.Context, cfg config.RabbitMQConfig) (*Client, error) { return client, nil } +// Config returns a copy of the effective rabbitmq configuration. Services that +// need to derive routing keys (e.g. desktop dispatch) consult this instead of +// hard-coding string constants. +func (c *Client) Config() config.RabbitMQConfig { + if c == nil { + return config.RabbitMQConfig{} + } + return c.cfg +} + func (c *Client) Close() error { if c == nil { return nil @@ -102,6 +112,91 @@ func (c *Client) PublishDesktopTaskEvent(ctx context.Context, body []byte) error return c.publishFanout(ctx, c.cfg.DesktopTaskEventExchange, body) } +func (c *Client) PublishDesktopDispatch(ctx context.Context, routingKey string, body []byte) error { + return c.publish(ctx, c.cfg.DesktopDispatchExchange, routingKey, body) +} + +// ConsumeDesktopDispatch declares a transient per-process queue bound to the +// desktop dispatch topic exchange using bindingKeys, and begins consuming from +// it. Every tenant-api instance gets its own queue so every instance sees every +// dispatch message; the hub inside the instance then filters by target client +// presence. Messages are auto-ack (transient delivery is OK: if the instance +// dies, the client reconnects and the server re-drives the task via the +// existing lease/recovery path). +func (c *Client) ConsumeDesktopDispatch(consumerName string, bindingKeys []string) (<-chan amqp.Delivery, *amqp.Channel, error) { + if len(bindingKeys) == 0 { + return nil, nil, fmt.Errorf("consume rabbitmq desktop dispatch: bindingKeys is required") + } + + for attempt := 0; attempt < 2; attempt++ { + conn, ch, err := c.openChannel() + if err != nil { + return nil, nil, err + } + + queue, err := ch.QueueDeclare( + "", + false, + true, + true, + false, + nil, + ) + if err != nil { + _ = ch.Close() + c.invalidateConnection(conn) + if attempt == 1 { + return nil, nil, fmt.Errorf("declare rabbitmq desktop dispatch queue: %w", err) + } + continue + } + + bindFailed := false + for _, key := range bindingKeys { + if err := ch.QueueBind( + queue.Name, + key, + c.cfg.DesktopDispatchExchange, + false, + nil, + ); err != nil { + _ = ch.Close() + c.invalidateConnection(conn) + if attempt == 1 { + return nil, nil, fmt.Errorf("bind rabbitmq desktop dispatch queue: %w", err) + } + bindFailed = true + break + } + } + if bindFailed { + continue + } + + deliveries, err := ch.Consume( + queue.Name, + consumerName, + true, + true, + false, + false, + nil, + ) + if err != nil { + _ = ch.Close() + c.invalidateConnection(conn) + if attempt == 1 { + return nil, nil, fmt.Errorf("consume rabbitmq desktop dispatch queue: %w", err) + } + continue + } + + return deliveries, ch, nil + } + + return nil, nil, fmt.Errorf("consume rabbitmq desktop dispatch queue: retry exhausted") +} + func (c *Client) PublishTemplateAssistTask(ctx context.Context, body []byte) error { return c.publish(ctx, c.cfg.TemplateAssistExchange, c.cfg.TemplateAssistRoutingKey, body) } @@ -578,6 +673,18 @@ func (c *Client) ensureTopology(conn *amqp.Connection) error { return fmt.Errorf("declare rabbitmq exchange %s: %w", c.cfg.DesktopTaskEventExchange, err) } + if err := ch.ExchangeDeclare( + c.cfg.DesktopDispatchExchange, + "topic", + true, + false, + false, + false, + nil, + ); err != nil { + return fmt.Errorf("declare rabbitmq exchange %s: %w", c.cfg.DesktopDispatchExchange, err) + } + if err := c.ensureWorkQueue( ch, c.cfg.TemplateAssistExchange, diff --git a/server/internal/shared/stream/desktop_dispatch.go b/server/internal/shared/stream/desktop_dispatch.go new file mode 100644 index 0000000..cd3bd7c --- /dev/null +++ b/server/internal/shared/stream/desktop_dispatch.go @@ -0,0 +1,274 @@ +package stream + +import ( + "context" + "encoding/json" + "fmt" + "os" + "strings" + "sync" + "time" + + amqp "github.com/rabbitmq/amqp091-go" + "go.uber.org/zap" + + "github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq" +) + +// DesktopDispatchEvent is the payload pushed to a desktop client over the +// dispatch WebSocket. The shape intentionally mirrors DesktopTaskEvent so that +// the existing SSE handler on the desktop client can treat both channels the +// same way and dedupe by task_id. +type DesktopDispatchEvent struct { + Type string `json:"type"` + TaskID string `json:"task_id"` + JobID string `json:"job_id,omitempty"` + WorkspaceID int64 `json:"workspace_id"` + TargetAccountID string `json:"target_account_id,omitempty"` + TargetClientID string `json:"target_client_id"` + Platform string `json:"platform,omitempty"` + Title *string `json:"title,omitempty"` + BusinessDate *string `json:"business_date,omitempty"` + SchedulerGroupKey *string `json:"scheduler_group_key,omitempty"` + QuestionText *string `json:"question_text,omitempty"` + Status string `json:"status"` + Kind string `json:"kind"` + Priority int `json:"priority,omitempty"` + UpdatedAt time.Time `json:"updated_at"` +} + +// DispatchSubscriber represents a live WebSocket connection bound to a single +// desktop client id. A client may have multiple subscribers (e.g. during a +// reconnect race), so the hub stores them as a set per client id. +type DispatchSubscriber struct { + id uint64 + out chan DesktopDispatchEvent + closed chan struct{} +} + +// Outbound returns the channel the handler goroutine should read from to write +// frames to the WebSocket. The channel is closed when the subscriber is +// unregistered; readers should range over it. +func (s *DispatchSubscriber) Outbound() <-chan DesktopDispatchEvent { + return s.out +} + +type dispatchClientStream struct { + subscribers map[uint64]*DispatchSubscriber +} + +// DesktopDispatchHub delivers per-client dispatch messages pulled from the +// desktop.task.dispatch topic exchange. Every tenant-api instance owns its own +// autodelete queue bound to the exchange, so every instance sees every message +// and routes it to the subscribers it currently owns. Clients not connected to +// this instance are simply filtered out; the client reconnects via the load +// balancer and lands on whichever instance currently holds its WebSocket. +type DesktopDispatchHub struct { + mu sync.RWMutex + clients map[string]*dispatchClientStream + rabbitMQ *rabbitmq.Client + logger *zap.Logger + instanceID string + consumerName string + bindingKeys []string + nextSubID uint64 +} + +// NewDesktopDispatchHub wires the hub to the shared rabbitmq client. The +// supplied binding keys decide which dispatch routing keys this instance +// receives; initially that is just `publish.*`, and the monitor slice will be +// added when the monitor chain is migrated. +func NewDesktopDispatchHub(rabbitMQClient *rabbitmq.Client, logger *zap.Logger, bindingKeys []string) *DesktopDispatchHub { + instanceID := fmt.Sprintf("%d-%d", os.Getpid(), time.Now().UnixNano()) + keys := make([]string, 0, len(bindingKeys)) + for _, key := range bindingKeys { + trimmed := strings.TrimSpace(key) + if trimmed == "" { + continue + } + keys = append(keys, trimmed) + } + if len(keys) == 0 { + keys = []string{"publish.*"} + } + return &DesktopDispatchHub{ + clients: make(map[string]*dispatchClientStream), + rabbitMQ: rabbitMQClient, + logger: logger, + instanceID: instanceID, + consumerName: "desktop-dispatch-" + instanceID, + bindingKeys: keys, + } +} + +// Run starts the background AMQP consumer. It is safe to call Run on a hub +// without a rabbitmq client — the hub simply behaves as an in-process broker +// used in tests, useful for unit-level wiring. +func (h *DesktopDispatchHub) Run(ctx context.Context) { + if h == nil || h.rabbitMQ == nil { + return + } + go h.run(ctx) +} + +// Subscribe registers a new WebSocket consumer for the given client id. The +// returned cancel closes the outbound channel and removes the subscriber; the +// WebSocket handler must call it on disconnect. +func (h *DesktopDispatchHub) Subscribe(clientID string) (*DispatchSubscriber, func()) { + if h == nil || clientID == "" { + return nil, func() {} + } + + h.mu.Lock() + stream, ok := h.clients[clientID] + if !ok { + stream = &dispatchClientStream{subscribers: make(map[uint64]*DispatchSubscriber)} + h.clients[clientID] = stream + } + h.nextSubID++ + sub := &DispatchSubscriber{ + id: h.nextSubID, + out: make(chan DesktopDispatchEvent, 32), + closed: make(chan struct{}), + } + stream.subscribers[sub.id] = sub + h.mu.Unlock() + + cancel := func() { + h.mu.Lock() + stream, ok := h.clients[clientID] + if !ok { + h.mu.Unlock() + return + } + existing, present := stream.subscribers[sub.id] + if !present { + h.mu.Unlock() + return + } + delete(stream.subscribers, sub.id) + if len(stream.subscribers) == 0 { + delete(h.clients, clientID) + } + h.mu.Unlock() + + select { + case <-existing.closed: + // already closed + default: + close(existing.closed) + close(existing.out) + } + } + + return sub, cancel +} + +// Publish delivers the event to every subscriber currently registered for the +// event's target client id on this instance. Events for clients that are not +// on this instance are dropped silently, by design. +func (h *DesktopDispatchHub) Publish(event DesktopDispatchEvent) { + if h == nil || event.TargetClientID == "" { + return + } + h.mu.RLock() + stream, ok := h.clients[event.TargetClientID] + if !ok { + h.mu.RUnlock() + return + } + subs := make([]*DispatchSubscriber, 0, len(stream.subscribers)) + for _, sub := range stream.subscribers { + subs = append(subs, sub) + } + h.mu.RUnlock() + + for _, sub := range subs { + select { + case <-sub.closed: + continue + case sub.out <- event: + default: + // Subscriber's outbound channel is full: the reader has fallen + // behind. We drop the message rather than block the dispatch + // consumer — the client will resync via the lease API on + // reconnect. + if h.logger != nil { + h.logger.Warn("desktop dispatch subscriber slow, dropping event", + zap.String("client_id", event.TargetClientID), + zap.String("task_id", event.TaskID), + ) + } + } + } +} + +// ActiveClientIDs returns the set of client ids currently connected on this +// instance. Intended for diagnostics and the /api/health/ready endpoint. +func (h *DesktopDispatchHub) ActiveClientIDs() []string { + if h == nil { + return nil + } + h.mu.RLock() + defer h.mu.RUnlock() + out := make([]string, 0, len(h.clients)) + for id := range h.clients { + out = append(out, id) + } + return out +} + +func (h *DesktopDispatchHub) run(ctx context.Context) { + for { + deliveries, ch, err := h.rabbitMQ.ConsumeDesktopDispatch(h.consumerName, h.bindingKeys) + if err != nil { + if h.logger != nil { + h.logger.Warn("desktop dispatch consume failed, retrying", + zap.Error(err), + ) + } + select { + case <-ctx.Done(): + return + case <-time.After(2 * time.Second): + continue + } + } + + closed := false + for !closed { + select { + case <-ctx.Done(): + _ = ch.Close() + return + case delivery, ok := <-deliveries: + if !ok { + closed = true + continue + } + h.handleDelivery(delivery) + } + } + _ = ch.Close() + } +} + +func (h *DesktopDispatchHub) handleDelivery(delivery amqp.Delivery) { + var event DesktopDispatchEvent + if err := json.Unmarshal(delivery.Body, &event); err != nil { + if h.logger != nil { + h.logger.Warn("desktop dispatch event decode failed", + zap.Error(err), + zap.String("routing_key", delivery.RoutingKey), + ) + } + return + } + if event.TargetClientID == "" { + // Fall back to the routing key: `{prefix}.{client_id}` -> client_id. + if idx := strings.LastIndex(delivery.RoutingKey, "."); idx >= 0 && idx < len(delivery.RoutingKey)-1 { + event.TargetClientID = delivery.RoutingKey[idx+1:] + } + } + h.Publish(event) +} diff --git a/server/internal/shared/stream/desktop_dispatch_test.go b/server/internal/shared/stream/desktop_dispatch_test.go new file mode 100644 index 0000000..ac62ffb --- /dev/null +++ b/server/internal/shared/stream/desktop_dispatch_test.go @@ -0,0 +1,135 @@ +package stream + +import ( + "testing" + "time" + + amqp "github.com/rabbitmq/amqp091-go" +) + +func TestDispatchHub_PublishRoutesToLocalSubscriber(t *testing.T) { + t.Parallel() + + hub := NewDesktopDispatchHub(nil, nil, []string{"publish.*"}) + sub, cancel := hub.Subscribe("client-a") + defer cancel() + + event := DesktopDispatchEvent{ + Type: "task_available", + TaskID: "task-1", + TargetClientID: "client-a", + Kind: "publish", + Status: "queued", + } + hub.Publish(event) + + select { + case got := <-sub.Outbound(): + if got.TaskID != "task-1" { + t.Fatalf("unexpected task id: %s", got.TaskID) + } + case <-time.After(200 * time.Millisecond): + t.Fatal("expected to receive event within 200ms") + } +} + +func TestDispatchHub_DropsEventsForUnknownClient(t *testing.T) { + t.Parallel() + + hub := NewDesktopDispatchHub(nil, nil, []string{"publish.*"}) + + // Subscribe for a different client, publish for another one. + _, cancel := hub.Subscribe("client-b") + defer cancel() + + hub.Publish(DesktopDispatchEvent{ + Type: "task_available", + TargetClientID: "client-a", + }) + + ids := hub.ActiveClientIDs() + if len(ids) != 1 || ids[0] != "client-b" { + t.Fatalf("expected only client-b subscribed, got %v", ids) + } +} + +func TestDispatchHub_CancelRemovesSubscriber(t *testing.T) { + t.Parallel() + + hub := NewDesktopDispatchHub(nil, nil, []string{"publish.*"}) + sub, cancel := hub.Subscribe("client-a") + + cancel() + + // Subsequent publish must not panic even though the subscriber is gone; + // channel is closed, reads must yield zero value with !ok. + hub.Publish(DesktopDispatchEvent{TargetClientID: "client-a"}) + + select { + case _, ok := <-sub.Outbound(): + if ok { + t.Fatal("expected outbound channel to be closed after cancel") + } + case <-time.After(50 * time.Millisecond): + t.Fatal("expected closed channel to yield immediately") + } + + if ids := hub.ActiveClientIDs(); len(ids) != 0 { + t.Fatalf("expected no subscribers after cancel, got %v", ids) + } +} + +func TestDispatchHub_MultipleSubscribersSameClient(t *testing.T) { + t.Parallel() + + hub := NewDesktopDispatchHub(nil, nil, []string{"publish.*"}) + subA, cancelA := hub.Subscribe("client-a") + defer cancelA() + subB, cancelB := hub.Subscribe("client-a") + defer cancelB() + + hub.Publish(DesktopDispatchEvent{ + Type: "task_available", + TaskID: "task-1", + TargetClientID: "client-a", + }) + + for i, ch := range []<-chan DesktopDispatchEvent{subA.Outbound(), subB.Outbound()} { + select { + case got := <-ch: + if got.TaskID != "task-1" { + t.Fatalf("subscriber %d got unexpected task id: %s", i, got.TaskID) + } + case <-time.After(200 * time.Millisecond): + t.Fatalf("subscriber %d did not receive event", i) + } + } +} + +func TestDispatchHub_HandleDeliveryFallsBackToRoutingKey(t *testing.T) { + t.Parallel() + + hub := NewDesktopDispatchHub(nil, nil, []string{"publish.*"}) + sub, cancel := hub.Subscribe("client-c") + defer cancel() + + // Body intentionally omits target_client_id; the hub should recover it + // from the routing key ".". + delivery := amqp.Delivery{ + RoutingKey: "publish.client-c", + Body: []byte(`{"type":"task_available","task_id":"task-9","kind":"publish","status":"queued"}`), + } + hub.handleDelivery(delivery) + + select { + case got := <-sub.Outbound(): + if got.TargetClientID != "client-c" { + t.Fatalf("expected fallback to set target client id, got %q", got.TargetClientID) + } + if got.TaskID != "task-9" { + t.Fatalf("unexpected task id: %s", got.TaskID) + } + case <-time.After(200 * time.Millisecond): + t.Fatal("expected fallback-routed event") + } +} diff --git a/server/internal/tenant/app/publish_job_service.go b/server/internal/tenant/app/publish_job_service.go index cc78b42..50615b2 100644 --- a/server/internal/tenant/app/publish_job_service.go +++ b/server/internal/tenant/app/publish_job_service.go @@ -2,7 +2,9 @@ package app import ( "context" + "encoding/json" "errors" + "fmt" "strings" "time" @@ -16,6 +18,7 @@ import ( sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache" "github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq" "github.com/geo-platform/tenant-api/internal/shared/response" + "github.com/geo-platform/tenant-api/internal/shared/stream" "github.com/geo-platform/tenant-api/internal/tenant/repository" ) @@ -232,6 +235,23 @@ func (s *PublishJobService) Create(ctx context.Context, actor auth.Actor, req Cr }); publishErr != nil { s.logWarn("desktop task available event publish failed", publishErr, zap.String("task_id", task.DesktopID.String())) } + + s.publishDispatchAvailable(task, stream.DesktopDispatchEvent{ + Type: "task_available", + TaskID: task.DesktopID.String(), + JobID: task.JobID.String(), + WorkspaceID: task.WorkspaceID, + TargetAccountID: task.TargetAccountID.String(), + TargetClientID: task.TargetClientID.String(), + Platform: task.Platform, + Title: title, + BusinessDate: businessDate, + SchedulerGroupKey: schedulerGroupKey, + QuestionText: questionText, + Status: task.Status, + Kind: task.Kind, + UpdatedAt: task.UpdatedAt, + }) } return &CreatePublishJobResponse{ @@ -351,6 +371,43 @@ func resolveRetryArticleIDFromPayload(payload map[string]any) (int64, bool) { return 0, false } +// publishDispatchAvailable pushes a per-client dispatch frame onto the topic +// exchange. Every tenant-api instance consumes the exchange and forwards the +// frame to the WebSocket owning the target client id. This is the primary +// signal the desktop client reacts to; the fanout event above is kept only so +// that other UIs (admin-web) see the same task lifecycle. +func (s *PublishJobService) publishDispatchAvailable(task *repository.DesktopTask, event stream.DesktopDispatchEvent) { + if s == nil || s.messaging == nil || task == nil { + return + } + + cfg := s.messaging.Config() + prefix := strings.TrimSpace(cfg.DesktopDispatchPublishPrefix) + if prefix == "" { + prefix = "publish" + } + targetClientID := task.TargetClientID.String() + if targetClientID == "" { + return + } + routingKey := fmt.Sprintf("%s.%s", prefix, targetClientID) + + body, err := json.Marshal(event) + if err != nil { + s.logWarn("desktop dispatch event encode failed", err, + zap.String("task_id", task.DesktopID.String()), + ) + return + } + + if err := s.messaging.PublishDesktopDispatch(context.Background(), routingKey, body); err != nil { + s.logWarn("desktop dispatch publish failed", err, + zap.String("task_id", task.DesktopID.String()), + zap.String("routing_key", routingKey), + ) + } +} + func (s *PublishJobService) logWarn(message string, err error, fields ...zap.Field) { if s.logger == nil || err == nil { return diff --git a/server/internal/tenant/transport/desktop_dispatch_handler.go b/server/internal/tenant/transport/desktop_dispatch_handler.go new file mode 100644 index 0000000..fc5e558 --- /dev/null +++ b/server/internal/tenant/transport/desktop_dispatch_handler.go @@ -0,0 +1,141 @@ +package transport + +import ( + "encoding/json" + "net/http" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "go.uber.org/zap" + + "github.com/geo-platform/tenant-api/internal/bootstrap" + "github.com/geo-platform/tenant-api/internal/shared/stream" +) + +const ( + dispatchWriteWait = 10 * time.Second + dispatchPongWait = 60 * time.Second + dispatchPingPeriod = 25 * time.Second + dispatchReadLimitBytes = 2048 + dispatchCloseGracePause = 500 * time.Millisecond +) + +type DesktopDispatchHandler struct { + app *bootstrap.App + upgrader websocket.Upgrader +} + +func NewDesktopDispatchHandler(a *bootstrap.App) *DesktopDispatchHandler { + return &DesktopDispatchHandler{ + app: a, + upgrader: websocket.Upgrader{ + ReadBufferSize: 4096, + WriteBufferSize: 4096, + CheckOrigin: func(_ *http.Request) bool { + // Origin enforcement happens at the gateway / LB layer; this + // endpoint is already gated by the bearer token middleware. + return true + }, + }, + } +} + +// Dispatch upgrades the HTTP request into a WebSocket and streams dispatch +// events for the current desktop client. The contract intentionally mirrors +// the existing SSE events endpoint: connect with a bearer token, receive JSON +// frames whose shape matches DesktopDispatchEvent, and reconnect on drop. +func (h *DesktopDispatchHandler) Dispatch(c *gin.Context) { + client := MustDesktopClient(c.Request.Context()) + hub := h.app.DesktopDispatch + if hub == nil { + c.AbortWithStatus(http.StatusServiceUnavailable) + return + } + + conn, err := h.upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + if h.app.Logger != nil { + h.app.Logger.Warn("desktop dispatch upgrade failed", + zap.String("client_id", client.ID.String()), + zap.Error(err), + ) + } + return + } + + subscriber, cancel := hub.Subscribe(client.ID.String()) + defer cancel() + + conn.SetReadLimit(dispatchReadLimitBytes) + _ = conn.SetReadDeadline(time.Now().Add(dispatchPongWait)) + conn.SetPongHandler(func(string) error { + _ = conn.SetReadDeadline(time.Now().Add(dispatchPongWait)) + return nil + }) + + // Reader goroutine: drains control frames and signals disconnect. We do + // not expect application data from the client over this socket. + readerDone := make(chan struct{}) + go func() { + defer close(readerDone) + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + }() + + connected := stream.DesktopDispatchEvent{ + Type: "connected", + TargetClientID: client.ID.String(), + WorkspaceID: client.WorkspaceID, + UpdatedAt: time.Now().UTC(), + } + if err := writeDispatchFrame(conn, connected); err != nil { + _ = conn.Close() + return + } + + pingTicker := time.NewTicker(dispatchPingPeriod) + defer pingTicker.Stop() + + defer func() { + _ = conn.SetWriteDeadline(time.Now().Add(dispatchWriteWait)) + _ = conn.WriteMessage(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) + time.Sleep(dispatchCloseGracePause) + _ = conn.Close() + }() + + outbound := subscriber.Outbound() + for { + select { + case <-c.Request.Context().Done(): + return + case <-readerDone: + return + case event, ok := <-outbound: + if !ok { + return + } + if err := writeDispatchFrame(conn, event); err != nil { + return + } + case <-pingTicker.C: + _ = conn.SetWriteDeadline(time.Now().Add(dispatchWriteWait)) + if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil { + return + } + } + } +} + +func writeDispatchFrame(conn *websocket.Conn, event stream.DesktopDispatchEvent) error { + body, err := json.Marshal(event) + if err != nil { + return err + } + _ = conn.SetWriteDeadline(time.Now().Add(dispatchWriteWait)) + return conn.WriteMessage(websocket.TextMessage, body) +} diff --git a/server/internal/tenant/transport/router.go b/server/internal/tenant/transport/router.go index c022f7f..b9bf3b9 100644 --- a/server/internal/tenant/transport/router.go +++ b/server/internal/tenant/transport/router.go @@ -29,6 +29,7 @@ func RegisterRoutes(app *bootstrap.App) { desktopAccountHandler := NewDesktopAccountHandler(app) desktopTaskHandler := NewDesktopTaskHandler(app) desktopEventsHandler := NewDesktopEventsHandler(app) + desktopDispatchHandler := NewDesktopDispatchHandler(app) desktopContentHandler := NewDesktopContentHandler(app) publishJobHandler := NewPublishJobHandler(app) protected.POST("/desktop/clients/register", desktopClientHandler.Register) @@ -40,6 +41,7 @@ func RegisterRoutes(app *bootstrap.App) { desktopAuth.POST("/clients/offline", desktopClientHandler.Offline) desktopAuth.POST("/clients/revoke", desktopClientHandler.Revoke) desktopAuth.GET("/events", desktopEventsHandler.Stream) + desktopAuth.GET("/dispatch", desktopDispatchHandler.Dispatch) desktopAuth.GET("/accounts", desktopAccountHandler.List) desktopAuth.GET("/content/articles/:id", desktopContentHandler.Article) desktopAuth.GET("/publish-tasks", desktopTaskHandler.ListPublish)