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"
	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"
)

// 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

	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"`

	// 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, TypeArtifactCreated:
		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.
func StartSession(workflowID string) []byte {
	return encode("StartSession", map[string]any{"workflowId": workflowID, "config": nil})
}

// 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,
	})
}