From 55a18b4b3a9b38db0d71c29927a496a6c8af5cf4 Mon Sep 17 00:00:00 2001 From: kami Date: Mon, 1 Jun 2026 13:36:01 +0400 Subject: [PATCH] =?UTF-8?q?feat:=20correx-managed=20model=20lifecycle=20sl?= =?UTF-8?q?ice=205=20=E2=80=94=20Go=20TUI=20model=20swap=20+=20telemetry?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Go TUI decodes model.changed / model.list / resource.status and renders a status-bar VRAM/GPU%/RAM gauge plus a models overlay (m key / palette 'models'): lists configured models with the resident one marked, enter swaps (SwapModel), c clears the pin (ClearModelPin). All three new server messages are non-event-bearing control/gauge frames, applied immediately like provider.status_changed. Server addition required for the picker: ServerMessage.ModelList (model.list, NonEventMessage) sent once in the initial snapshot from ManagedInferenceRouter.availableModelIds()/currentModelId() — the TUI otherwise cannot enumerate swap targets. Completes the 5-slice model-lifecycle feature. ./gradlew check green; go build/vet clean. Plan: docs/plans/2026-05-31-model-lifecycle-management.md (slice 5 of 5). --- .../apps/server/protocol/ServerMessage.kt | 13 +++++++ .../apps/server/ws/GlobalStreamHandler.kt | 10 +++++ apps/tui-go/internal/app/model.go | 9 +++++ apps/tui-go/internal/app/overlays.go | 34 +++++++++++++++++ apps/tui-go/internal/app/server.go | 19 +++++++++- apps/tui-go/internal/app/update.go | 38 +++++++++++++++++++ apps/tui-go/internal/app/view.go | 20 +++++++++- apps/tui-go/internal/protocol/protocol.go | 27 +++++++++++++ .../commons/ManagedInferenceRouter.kt | 6 +++ 9 files changed, 174 insertions(+), 2 deletions(-) diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ServerMessage.kt b/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ServerMessage.kt index 1a184f98..37c7f3fc 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ServerMessage.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ServerMessage.kt @@ -283,6 +283,19 @@ sealed interface ServerMessage { override val sessionSequence: Long? = null, ) : ServerMessage, NonEventMessage + /** + * The configured local models and which one is currently resident. Sent once in the initial + * snapshot when correx manages the model process, so clients can offer a swap picker. + */ + @Serializable + @SerialName("model.list") + data class ModelList( + val models: List, + val current: String?, + override val sequence: Long? = null, + override val sessionSequence: Long? = null, + ) : ServerMessage, NonEventMessage + /** * Live resource gauge pushed periodically on the global stream (not event-derived). All fields * are nullable: GPU fields are null on a non-NVIDIA host, `processRssMb` is null when no model diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/ws/GlobalStreamHandler.kt b/apps/server/src/main/kotlin/com/correx/apps/server/ws/GlobalStreamHandler.kt index 1a7053e5..7a77bbe3 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/ws/GlobalStreamHandler.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/ws/GlobalStreamHandler.kt @@ -149,6 +149,16 @@ class GlobalStreamHandler(private val module: ServerModule) { ServerMessage.WorkflowList(workflows = workflows), ))) + // Configured models + current, so the client can offer a swap picker (managed path only). + module.modelSwapper?.let { swapper -> + session.send(Frame.Text(ProtocolSerializer.encodeServerMessage( + ServerMessage.ModelList( + models = swapper.availableModelIds(), + current = swapper.currentModelId(), + ), + ))) + } + // Initial resource gauge so a freshly-connected client renders VRAM/RAM immediately. val snapshot = withContext(Dispatchers.IO) { module.resourceProbe.probe() } session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(snapshot.toResourceStatus()))) diff --git a/apps/tui-go/internal/app/model.go b/apps/tui-go/internal/app/model.go index c476e72e..ac30e1c5 100644 --- a/apps/tui-go/internal/app/model.go +++ b/apps/tui-go/internal/app/model.go @@ -45,6 +45,7 @@ const ( OverlayEventInspector OverlayDiff OverlayToolPalette + OverlayModels ) // RouterEntry is one line in a session's conversation transcript. @@ -158,6 +159,14 @@ type Model struct { currentModel string providerType string // LOCAL | REMOTE + // managed-model swap (nil resource fields = unavailable) + availableModels []string + modelsIndex int + gpuUsedMB *int64 + gpuTotalMB *int64 + gpuUtil *int + ramMB *int64 + // diff / overlay overlay OverlayKind overlayEventIdx int diff --git a/apps/tui-go/internal/app/overlays.go b/apps/tui-go/internal/app/overlays.go index bac13e5d..bfc24b95 100644 --- a/apps/tui-go/internal/app/overlays.go +++ b/apps/tui-go/internal/app/overlays.go @@ -37,6 +37,8 @@ func (m Model) renderOverlay(base string) string { return m.center(m.diffModal()) case OverlayToolPalette: return m.center(m.toolPaletteModal()) + case OverlayModels: + return m.center(m.modelsModal()) } return base } @@ -243,6 +245,38 @@ func (m Model) toolPaletteModal() string { return m.center(modal) } +func (m Model) modelsModal() string { + t := m.theme + w := m.modalWidth() + + var b strings.Builder + b.WriteString(m.titleLine("models") + "\n\n") + + if g := m.gaugeText(); g != "" { + b.WriteString(mbg(t, g, t.P.Faint) + "\n\n") + } + + if len(m.availableModels) == 0 { + b.WriteString(mbg(t, " model management not enabled", t.P.Faint) + "\n") + } + for i, id := range m.availableModels { + marker := mbg(t, " ", t.P.BgPanel) + nameFg := t.P.Fg + if i == m.modelsIndex { + marker = lipgloss.NewStyle().Foreground(t.P.Accent).Background(t.P.BgPanel).Render("▌ ") + nameFg = t.P.FgStrong + } + line := marker + lipgloss.NewStyle().Foreground(nameFg).Background(t.P.BgPanel).Render(padRaw(id, 28)) + if id == m.currentModel { + line += lipgloss.NewStyle().Foreground(t.P.Accent2).Background(t.P.BgPanel).Render(" ◂ resident") + } + b.WriteString(line + "\n") + } + b.WriteString("\n" + modalHints(t, [][2]string{{"↑↓", "select"}, {"enter", "swap"}, {"c", "clear pin"}, {"esc", "close"}})) + modal := t.Overlay.Width(w).Render(b.String()) + return m.center(modal) +} + func tierColor(t Theme, tier int) lipgloss.Color { switch { case tier >= 3: diff --git a/apps/tui-go/internal/app/server.go b/apps/tui-go/internal/app/server.go index 55f77c8e..5b607181 100644 --- a/apps/tui-go/internal/app/server.go +++ b/apps/tui-go/internal/app/server.go @@ -190,6 +190,22 @@ func (m *Model) applyServer(msg protocol.ServerMessage) { for _, w := range msg.Workflows { m.workflows = append(m.workflows, Workflow{ID: w.WorkflowID, Description: w.Description}) } + case protocol.TypeModelChanged: + if msg.Loaded { + m.currentModel = msg.ModelID + m.providerType = "LOCAL" + } + case protocol.TypeModelList: + m.availableModels = append(m.availableModels[:0], msg.Models...) + if msg.Current != "" { + m.currentModel = msg.Current + m.providerType = "LOCAL" + } + case protocol.TypeResourceStatus: + m.gpuUsedMB = msg.GpuMemoryUsedMb + m.gpuTotalMB = msg.GpuMemoryTotalMb + m.gpuUtil = msg.GpuUtilizationPct + m.ramMB = msg.ProcessRssMb } // Background-update badge for non-selected sessions. @@ -202,7 +218,8 @@ func sessionIDOf(msg protocol.ServerMessage) string { switch msg.Type { case protocol.TypeStageManifest, protocol.TypeSnapshotComplete, protocol.TypeProtocolError, protocol.TypeProviderStatus, - protocol.TypeWorkflowList, protocol.TypeRouterResponse: + protocol.TypeWorkflowList, protocol.TypeRouterResponse, + protocol.TypeModelChanged, protocol.TypeModelList, protocol.TypeResourceStatus: return "" default: return msg.SessionID diff --git a/apps/tui-go/internal/app/update.go b/apps/tui-go/internal/app/update.go index 77a94481..224c92f3 100644 --- a/apps/tui-go/internal/app/update.go +++ b/apps/tui-go/internal/app/update.go @@ -144,6 +144,8 @@ func (m Model) handleNormalKey(k tea.KeyMsg) (tea.Model, tea.Cmd) { m.overlayEventIdx = 0 case "t": m.overlay = OverlayToolPalette + case "m": + m.openModelsOverlay() case "w": if ds == StateIdle { m.wfVisible = !m.wfVisible @@ -333,10 +335,43 @@ func (m Model) handleOverlayKey(k tea.KeyMsg) (tea.Model, tea.Cmd) { if runeIs(k, "t") { m.overlay = OverlayNone } + case OverlayModels: + switch { + case k.Type == tea.KeyUp || runeIs(k, "k"): + if m.modelsIndex > 0 { + m.modelsIndex-- + } + case k.Type == tea.KeyDown || runeIs(k, "j"): + if m.modelsIndex < len(m.availableModels)-1 { + m.modelsIndex++ + } + case k.Type == tea.KeyEnter: + if m.modelsIndex >= 0 && m.modelsIndex < len(m.availableModels) { + m.client.Send(protocol.SwapModel(m.availableModels[m.modelsIndex])) + m.overlay = OverlayNone + } + case runeIs(k, "c"): + m.client.Send(protocol.ClearModelPin()) + m.overlay = OverlayNone + case runeIs(k, "m"): + m.overlay = OverlayNone + } } return m, nil } +// openModelsOverlay opens the model picker, pre-selecting the resident model. +func (m *Model) openModelsOverlay() { + m.overlay = OverlayModels + m.modelsIndex = 0 + for i, id := range m.availableModels { + if id == m.currentModel { + m.modelsIndex = i + break + } + } +} + func runeIs(k tea.KeyMsg, s string) bool { return k.Type == tea.KeyRunes && string(k.Runes) == s } @@ -374,6 +409,7 @@ func paletteCommands() []paletteCmd { return []paletteCmd{ {"workflows", "start workflow", "open the workflow picker"}, {"tools", "tool palette", "tools for the current stage"}, + {"models", "swap model", "pick / pin the local model"}, {"events", "event inspector", "browse the event stream"}, {"mode", "toggle mode", "switch chat / steering"}, {"cancel", "cancel session", "stop the selected session"}, @@ -404,6 +440,8 @@ func (m Model) execPalette(id string) (tea.Model, tea.Cmd) { m.wfIndex = 0 case "tools": m.overlay = OverlayToolPalette + case "models": + m.openModelsOverlay() case "events": m.overlay = OverlayEventInspector m.overlayEventIdx = 0 diff --git a/apps/tui-go/internal/app/view.go b/apps/tui-go/internal/app/view.go index 976b97b5..6818a07e 100644 --- a/apps/tui-go/internal/app/view.go +++ b/apps/tui-go/internal/app/view.go @@ -81,6 +81,9 @@ func (m Model) renderStatus() string { left := strings.Join(parts, sep) var right []string + if g := m.gaugeText(); g != "" { + right = append(right, span(g, t.P.Faint)) + } if s := m.session(m.selectedID); s != nil && s.Active { right = append(right, lipgloss.NewStyle().Foreground(t.P.Accent).Background(bg).Render(m.spinner())+span(" active", t.P.Accent2)) } @@ -93,6 +96,21 @@ func (m Model) renderStatus() string { return m.justify(left, strings.Join(right, span(" · ", t.P.Faint)), m.width, bg) } +// gaugeText renders the live VRAM/GPU/RAM gauge, omitting any unavailable field. +func (m Model) gaugeText() string { + var parts []string + if m.gpuUsedMB != nil && m.gpuTotalMB != nil { + parts = append(parts, "VRAM "+itoa(int(*m.gpuUsedMB))+"/"+itoa(int(*m.gpuTotalMB))+"M") + } + if m.gpuUtil != nil { + parts = append(parts, "GPU "+itoa(*m.gpuUtil)+"%") + } + if m.ramMB != nil { + parts = append(parts, "RAM "+itoa(int(*m.ramMB))+"M") + } + return strings.Join(parts, " ") +} + var spinnerFrames = []string{"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"} func (m Model) spinner() string { return spinnerFrames[m.frame%len(spinnerFrames)] } @@ -121,7 +139,7 @@ func (m Model) renderFooter() string { case m.displayState() == StateIdle: hints = []string{hint("i", "name"), hint("/", "filter"), hint("enter", "open"), hint("jk", "move"), hint("w", "workflows"), hint("p", "cmds"), hint("q", "quit")} default: // StateInSession - hints = []string{hint("i", "message"), hint("e", "events"), hint("t", "tools"), hint("s", "mode"), hint("l", "back"), hint("p", "cmds"), hint("q", "quit")} + hints = []string{hint("i", "message"), hint("e", "events"), hint("t", "tools"), hint("m", "model"), hint("s", "mode"), hint("l", "back"), hint("p", "cmds"), hint("q", "quit")} } left := strings.Join(hints, gap) diff --git a/apps/tui-go/internal/protocol/protocol.go b/apps/tui-go/internal/protocol/protocol.go index 25e909c1..19e3be16 100644 --- a/apps/tui-go/internal/protocol/protocol.go +++ b/apps/tui-go/internal/protocol/protocol.go @@ -38,6 +38,9 @@ const ( TypeSnapshotComplete = "snapshot_complete" TypeWorkflowList = "workflow.list" TypeToolAssessed = "tool.assessed" + TypeModelChanged = "model.changed" + TypeModelList = "model.list" + TypeResourceStatus = "resource.status" ) // ServerMessage is a flat decode of every server->client variant. Field names @@ -69,6 +72,20 @@ type ServerMessage struct { ElapsedMs int64 `json:"elapsedMs"` LastSequence int64 `json:"lastSequence"` + // model.changed + ModelID string `json:"modelId"` + Loaded bool `json:"loaded"` + + // model.list (current reuses the field below) + Models []string `json:"models"` + Current string `json:"current"` + + // resource.status (nil = unavailable) + GpuMemoryUsedMb *int64 `json:"gpuMemoryUsedMb"` + GpuMemoryTotalMb *int64 `json:"gpuMemoryTotalMb"` + GpuUtilizationPct *int `json:"gpuUtilizationPct"` + ProcessRssMb *int64 `json:"processRssMb"` + State *SessionStateDto `json:"state"` RiskSummary *RiskSummaryDto `json:"riskSummary"` Status *ProviderHealthDto `json:"status"` @@ -206,6 +223,16 @@ func Ping(ts int64) []byte { return encode("Ping", map[string]any{"timestamp": ts}) } +// SwapModel asks the server to make modelID resident and pin it. +func SwapModel(modelID string) []byte { + return encode("SwapModel", map[string]any{"modelId": modelID}) +} + +// ClearModelPin releases the manual-swap pin so per-stage selection resumes. +func ClearModelPin() []byte { + return encode("ClearModelPin", map[string]any{}) +} + // CreateGrant pre-authorizes a tool/tier so future calls skip the approval gate. // scope is "SESSION" or "STAGE"; permittedTiers are tier names like "T3". func CreateGrant(sessionID, scope string, permittedTiers []string, reason, toolName string) []byte { diff --git a/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt index ae1e3123..67110612 100644 --- a/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt +++ b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt @@ -79,6 +79,12 @@ class ManagedInferenceRouter( /** The currently pinned model id, or null when no manual swap is active. */ fun pinnedModel(): String? = pinnedModelId + /** All configured model ids, in config order — the swap candidates surfaced to clients. */ + fun availableModelIds(): List = ordered.map { it.modelId } + + /** The currently resident model id, or null when none is loaded. */ + fun currentModelId(): String? = manager.currentModel()?.modelId + /** * Selects the configured model whose declared capabilities cover every required capability, * preferring the highest summed score over those required capabilities. Returns null when no