correx/apps/tui-go/internal/protocol/protocol.go

// Package protocol implements the correx WebSocket wire contract used to talk
// to apps/server over /stream. The server (Kotlin/kotlinx.serialization) uses a
// "type" class discriminator. Server->client variants carry @SerialName values
// (e.g. "session.announced"); client->server variants have no @SerialName, so the
// discriminator is the fully-qualified Kotlin class name.
package protocol

import "encoding/json"

// clientPrefix is the kotlinx default serial name for nested ClientMessage
// subclasses (qualified class name). The server decodes against this exact value.
const clientPrefix = "com.correx.apps.server.protocol.ClientMessage."

// Server message discriminator values (@SerialName on ServerMessage variants).
const (
	TypeSessionAnnounced = "session.announced"
	TypeChatTurn         = "chat.turn"
	TypeSessionPaused    = "session.paused"
	TypeSessionResumed   = "session.resumed"
	TypeSessionCompleted = "session.completed"
	TypeSessionFailed    = "session.failed"
	TypeSessionSnapshot  = "session_snapshot"
	TypeStageStarted     = "stage.started"
	TypeStageCompleted   = "stage.completed"
	TypeStageFailed      = "stage.failed"
	TypeInferenceStarted = "inference.started"
	TypeInferenceDone    = "inference.completed"
	TypeInferenceTimeout = "inference.timed_out"
	TypeToolStarted      = "tool.started"
	TypeToolCompleted    = "tool.completed"
	TypeToolFailed       = "tool.failed"
	TypeToolRejected     = "tool.rejected"
	TypeApprovalRequired = "approval.required"
	TypeApprovalResolved = "approval.resolved"
	TypeWorkspaceBound   = "session.workspace_bound"
	TypeStageManifest    = "stage.tool_manifest"
	TypeProviderStatus   = "provider.status_changed"
	TypeProtocolError    = "protocol_error"
	TypeArtifactCreated  = "artifact.created"
	TypeRouterResponse   = "router.response"
	TypeSnapshotComplete = "snapshot_complete"
	TypeWorkflowList     = "workflow.list"
	TypeToolAssessed     = "tool.assessed"
	TypeModelChanged     = "model.changed"
	TypeModelList        = "model.list"
	TypeResourceStatus   = "resource.status"
	TypeRouterNarration  = "router.narration"
)

// ServerMessage is a flat decode of every server->client variant. Field names
// across variants don't collide in incompatible JSON types, so a single struct
// keyed by Type decodes the whole protocol. Value-class IDs serialize as plain
// strings (verified against ServerMessageSerializationTest).
type ServerMessage struct {
	Type string `json:"type"`

	SessionID     string  `json:"sessionId"`
	WorkflowID    string  `json:"workflowId"`
	StageID       string  `json:"stageId"`
	ToolName      string  `json:"toolName"`
	Role          string  `json:"role"`
	TurnID        string  `json:"turnId"`
	Reason        string  `json:"reason"`
	Summary       string  `json:"outputSummary"`
	Response      string  `json:"responseText"`
	Content       string  `json:"content"`
	Diff          *string `json:"diff"`
	Tier          string  `json:"tier"`
	Message       string  `json:"message"`
	RequestID     string  `json:"requestId"`
	ArtifactID    string  `json:"artifactId"`
	ProviderID    string  `json:"providerId"`
	Preview       *string `json:"preview"`
	Disposition   string  `json:"disposition"`
	Outcome       string  `json:"outcome"`       // approval.resolved: APPROVED | REJECTED | AUTO_APPROVED
	WorkspaceRoot string  `json:"workspaceRoot"` // session.workspace_bound: the bound cwd

	SteeringEmitted bool  `json:"steeringEmitted"`
	OccurredAt      int64 `json:"occurredAt"`
	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"`

	// chat.turn router metrics (nil when not a ROUTER turn or server omits them)
	LatencyMs   *int64 `json:"latencyMs"`
	TotalTokens *int   `json:"totalTokens"`

	// 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"`
	PendingAppr    []ApprovalDto      `json:"pendingApprovals"`
	Tools          []ToolRecordDto    `json:"tools"`
	RecentEvents   []EventEntryDto    `json:"recentEvents"`
	Stages         []StageToolDecl    `json:"stages"`
	Workflows      []WorkflowDto      `json:"workflows"`
	AssessedIssues []AssessedIssueDto `json:"issues"`
}

type SessionStateDto struct {
	Status         string  `json:"status"`
	CurrentStageID *string `json:"currentStageId"`
	PauseReason    *string `json:"pauseReason"`
}

type RiskSummaryDto struct {
	Level             string   `json:"level"`
	Factors           []string `json:"factors"`
	RecommendedAction string   `json:"recommendedAction"`
	Rationale         []string `json:"rationale"`
}

type AssessedIssueDto struct {
	Code     string `json:"code"`
	Message  string `json:"message"`
	Severity string `json:"severity"`
}

type ProviderHealthDto struct {
	ProviderID  string `json:"providerId"`
	Status      string `json:"status"`
	LoadPercent *int   `json:"loadPercent"`
}

type ApprovalDto struct {
	RequestID string  `json:"requestId"`
	Tier      string  `json:"tier"`
	ToolName  *string `json:"toolName"`
	Preview   *string `json:"preview"`
}

type StageToolDecl struct {
	StageID string     `json:"stageId"`
	Tools   []ToolDecl `json:"tools"`
}

type ToolDecl struct {
	Name string `json:"name"`
	Tier int    `json:"tier"`
}

type EventEntryDto struct {
	Timestamp int64  `json:"timestamp"`
	Type      string `json:"type"`
	Detail    string `json:"detail"`
}

type ToolRecordDto struct {
	Name   string  `json:"name"`
	Tier   int     `json:"tier"`
	Status string  `json:"status"`
	Diff   *string `json:"diff"`
}

type WorkflowDto struct {
	WorkflowID  string `json:"workflowId"`
	Description string `json:"description"`
}

// Decode parses one server frame into a ServerMessage.
func Decode(raw []byte) (ServerMessage, error) {
	var m ServerMessage
	err := json.Unmarshal(raw, &m)
	return m, err
}

// IsEventBearing reports whether a message carries a session event (and so must
// be buffered during the snapshot phase). Control/snapshot messages return false
// and are applied immediately.
func (m ServerMessage) IsEventBearing() bool {
	switch m.Type {
	case TypeSessionAnnounced, TypeChatTurn, TypeSessionPaused, TypeSessionResumed,
		TypeSessionCompleted, TypeSessionFailed,
		TypeStageStarted, TypeStageCompleted, TypeStageFailed,
		TypeInferenceStarted, TypeInferenceDone, TypeInferenceTimeout,
		TypeToolStarted, TypeToolCompleted, TypeToolFailed, TypeToolRejected,
		TypeToolAssessed,
		TypeApprovalRequired, TypeApprovalResolved, TypeWorkspaceBound, TypeArtifactCreated,
		TypeRouterNarration:
		return true
	default:
		return false
	}
}

// --- Client -> server encoders ---

func encode(typeName string, fields map[string]any) []byte {
	fields["type"] = clientPrefix + typeName
	b, _ := json.Marshal(fields)
	return b
}

// StartChatSession opens a new chat session with a client-generated sessionId.
func StartChatSession(sessionID, text string) []byte {
	return encode("StartChatSession", map[string]any{"sessionId": sessionID, "text": text})
}

// ChatInput sends a turn into an existing session. mode is "CHAT" or "STEERING".
func ChatInput(sessionID, text, mode string) []byte {
	return encode("ChatInput", map[string]any{"sessionId": sessionID, "text": text, "mode": mode})
}

// StartSession launches a workflow-backed session. input is the optional freeform request
// seeded into the decision journal; pass "" for fixed-task workflows (e.g. healthcheck).
func StartSession(workflowID, input string) []byte {
	msg := map[string]any{"workflowId": workflowID, "config": nil}
	if input != "" {
		msg["input"] = input
	}
	return encode("StartSession", msg)
}

// CancelSession cancels a running session.
func CancelSession(sessionID string) []byte {
	return encode("CancelSession", map[string]any{"sessionId": sessionID})
}

// ApprovalResponse answers an approval gate. decision is "APPROVE" or "REJECT".
func ApprovalResponse(requestID, decision string, steeringNote *string) []byte {
	return encode("ApprovalResponse", map[string]any{
		"requestId":    requestID,
		"decision":     decision,
		"steeringNote": steeringNote,
	})
}

// Ping is a keepalive.
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 {
	return encode("CreateGrant", map[string]any{
		"sessionId":      sessionID,
		"scope":          scope,
		"stageId":        nil,
		"permittedTiers": permittedTiers,
		"reason":         reason,
		"expiresAt":      nil,
		"toolName":       toolName,
	})
}

// Hello is the first frame sent on every (re)connect. It carries the client's
// working directory so the server can bind a workspace before any session starts.
func Hello(workingDir string) []byte {
	return encode("Hello", map[string]any{"workingDir": workingDir})
}