Cyrene_v2/internal/services/gateway/ws_route.go

package gateway

import (
	"context"
	"encoding/json"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/auth"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/orchestration"
)

// wsForbiddenCode is the machine-readable code of the frame a non owner/admin
// connection gets for any inbound instruction (design/11 §7.2 B). It rides the
// file's existing error-frame shape — wsEvent{Type: "error", Data: map[string]any}
// — so the client decodes it with the code it already knows, and distinguishes
// 「你没有这个权限」 from a handler-level failure by the key alone.
const wsForbiddenCode = "forbidden"

// route dispatches one decoded inbound (client→server) frame by its type.
//
// Role gate (design/11 §7.2 B, decided by the owner on 2026-10-03): a connection
// that is not owner/admin may not send ANY in-app instruction. The check is the
// first statement here and it denies by default — it is deliberately NOT a list of
// forbidden types — so a frame type added by a later milestone is refused until
// someone registers it on purpose. A viewer's /ws/chat connection is therefore a
// read-only monitor: it follows the session state machine but can never drive the
// model, cancel a turn, read a sub-agent's context or a task's output, or feed
// voice input. Each refused frame is answered on that one connection with
// {type:"error", data:{code:"forbidden"}}; nothing is broadcast, no handler runs,
// so no model call is made and nothing is persisted.
//
// Widening this: a viewer-side inbound frame may be allowed only after it is given
// an opening and registered in design/11 §7.2 (「必须先在设计里给它开口并登记白名单」).
// The blanket deny below is what makes forgetting that step safe, so allowing one
// means replacing it with a registry, not deleting it.
//
// Protocol-level frames (ping/pong/close) never reach here — gorilla handles them
// inside ReadMessage — so the gate cannot kill a connection. A frame that is not
// JSON at all is answered by readPump before dispatch; it carries no type and can
// trigger nothing, which is why it keeps its parsing-level answer.
func (c *wsConn) route(in wsInbound) {
	if !ownerAdminClaims(c.claims) {
		c.h.log.Debug("ws inbound rejected", "type", in.Type, "role", claimRole(c.claims))
		c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"code": wsForbiddenCode}})
		return
	}
	switch in.Type {
	case "mode_switch":
		c.publishPublic("ws.mode_switch", rawData(in.Data))
	case "chat":
		c.publishPublic(wsChatSubject, rawData(in.Data))
	case "interrupt":
		// design/15 §3: user says "停下/算了" → cancel the in-flight turn and
		// clear the session's pending queue, so the next message starts fresh.
		c.h.handleInterrupt(in.Data)
	case "voice_input":
		// Voice routes through the injected STT→chat→TTS pipeline. Without a
		// voice service it replies with a clear error rather than dropping the
		// frame. The audio reply is delivered as a voice_output event on the
		// requesting connection.
		c.handleVoiceInput(in)
	case "subagent_ask":
		c.handleSubagentAsk(in)
	case "task_output_query":
		c.handleTaskOutputQuery(in)
	default:
		c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"message": "unknown event type"}})
	}
}

// subagentAskRequest is the decoded payload of a subagent_ask frame. Mode selects
// the audit window (""/full/tail/keyword); n bounds a tail query; query is the
// keyword for a keyword query.
type subagentAskRequest struct {
	ID    string `json:"id"`
	Mode  string `json:"mode"`
	N     int    `json:"n"`
	Query string `json:"query"`
}

// taskOutputRequest is the decoded payload of a task_output_query frame.
type taskOutputRequest struct {
	TaskID string `json:"taskId"`
}

// handleSubagentAsk audits a sub-agent's context and replies with a
// ws.subagent_event frame carrying the real context window. When no auditor is
// injected (gateway not wired with a scheduler) it replies with a clear error
// rather than a fabricated "pending" snapshot.
func (c *wsConn) handleSubagentAsk(in wsInbound) {
	if c.h.auditor == nil {
		c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"message": "subagent audit unavailable"}})
		return
	}
	var req subagentAskRequest
	_ = json.Unmarshal(in.Data, &req)
	if req.ID == "" {
		c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"message": "subagent id required"}})
		return
	}
	mode := orchestration.AuditMode(req.Mode)
	if mode == "" {
		mode = orchestration.AuditFull
	}
	res, err := c.h.auditor.AuditContext(context.Background(), req.ID, mode, req.N, req.Query)
	if err != nil {
		c.h.log.Debug("subagent audit", "id", req.ID, "err", err)
		c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"message": err.Error()}})
		return
	}
	c.enqueueJSON(wsEvent{Type: "subagent_event", Data: res})
}

// handleTaskOutputQuery fetches the real task output for taskId and replies with
// a ws.task_output frame. When no task store is injected it replies with a clear
// error rather than an empty {}. A task that is still running returns the output
// captured so far.
func (c *wsConn) handleTaskOutputQuery(in wsInbound) {
	if c.h.taskOutputs == nil {
		c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"message": "task output unavailable"}})
		return
	}
	var req taskOutputRequest
	_ = json.Unmarshal(in.Data, &req)
	if req.TaskID == "" {
		c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"message": "task id required"}})
		return
	}
	ev, err := c.h.taskOutputs.TaskOutput(context.Background(), req.TaskID)
	if err != nil {
		c.h.log.Debug("task output query", "id", req.TaskID, "err", err)
		c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"message": err.Error()}})
		return
	}
	text := ev.Output.Stdout
	if ev.Output.Stderr != "" {
		if text != "" {
			text += "\n"
		}
		text += ev.Output.Stderr
	}
	c.enqueueJSON(wsEvent{Type: "task_output", Data: map[string]any{
		"taskId": ev.ID,
		"output": text,
		"status": string(ev.Status),
	}})
}

// claimRole labels a connection's role for the inbound-rejection log record. A
// connection with nil claims (the fail-closed bare connection a unit test builds)
// logs as "none" instead of dereferencing the nil pointer.
func claimRole(claims *auth.Claims) string {
	if claims == nil {
		return "none"
	}
	return string(claims.Role)
}