server.go

package relay

import "context"

// -----------------------------------------------------------------------------
// Agent Registry Interface
// -----------------------------------------------------------------------------

// AgentRegistry manages agent connections and provides agent discovery.
// Implementations can store agent state in memory, Redis, or other backends.
type AgentRegistry interface {
	// RegisterAgent registers a new agent connection with the registry.
	// Called by the transport layer when an agent connects.
	RegisterAgent(ctx context.Context, conn AgentConn) error

	// UnregisterAgent removes an agent from the registry.
	// Called by the transport layer when an agent disconnects.
	UnregisterAgent(ctx context.Context, id AgentID) error

	// GetAgent returns information about a specific agent.
	// Returns an error if the agent is not connected.
	GetAgent(ctx context.Context, id AgentID) (*AgentInfo, error)

	// GetAgentConn returns the connection for a specific agent.
	// Used to send messages to the agent.
	GetAgentConn(ctx context.Context, id AgentID) (AgentConn, error)

	// ListAgents returns information about all connected agents.
	// Can be filtered by capabilities in the options.
	ListAgents(ctx context.Context, opts *ListAgentsOptions) ([]*AgentInfo, error)

	// HandleHeartbeat processes a heartbeat from an agent.
	// Updates the agent's last heartbeat timestamp.
	HandleHeartbeat(ctx context.Context, agentID AgentID) error
}

// ListAgentsOptions provides filtering options for ListAgents.
type ListAgentsOptions struct {
	// Labels filters agents by label key-value pairs (all must match).
	Labels map[string]string

	// Capabilities filters agents by required capabilities (all must be present).
	Capabilities []Capability
}

// -----------------------------------------------------------------------------
// Client Registry Interface
// -----------------------------------------------------------------------------

// ClientRegistry manages client connections.
// Implementations can store client state in memory, Redis, or other backends.
// This mirrors the AgentRegistry pattern to enable symmetric routing.
type ClientRegistry interface {
	// RegisterClient registers a new client connection with the registry.
	// Called by the transport layer when a client connects.
	RegisterClient(ctx context.Context, conn ClientConn) error

	// UnregisterClient removes a client from the registry.
	// Called by the transport layer when a client disconnects.
	UnregisterClient(ctx context.Context, id ClientID) error

	// GetClientConn returns the connection for a specific client.
	// Used to send messages to the client.
	GetClientConn(ctx context.Context, id ClientID) (ClientConn, error)
}

// -----------------------------------------------------------------------------
// Session Manager Interface
// -----------------------------------------------------------------------------

// SessionManager manages session lifecycle and message routing.
// Implementations handle session state and coordinate with agents.
type SessionManager interface {
	// OpenSession opens a new session on the specified agent for a client.
	// The session type and options determine what kind of session is created.
	// Returns the session ID that can be used to interact with the session.
	OpenSession(ctx context.Context, clientID ClientID, agentID AgentID, sessionType SessionType, opts *SessionOptions) (SessionID, error)

	// CloseSession terminates a session and releases resources.
	CloseSession(ctx context.Context, sessionID SessionID) error

	// GetSession returns the session interface for direct manipulation.
	// Returns an error if the session doesn't exist.
	GetSession(ctx context.Context, sessionID SessionID) (Session, error)

	// HandleSessionMessage processes a session-related message from an agent.
	// Routes the message to the appropriate session handler.
	HandleSessionMessage(ctx context.Context, agentID AgentID, msg *AgentMessage) error

	// CloseAgentSessions closes all sessions for a specific agent.
	// Called when an agent disconnects.
	CloseAgentSessions(ctx context.Context, agentID AgentID) error

	// CloseClientSessions closes all sessions for a specific client.
	// Called when a client disconnects.
	CloseClientSessions(ctx context.Context, clientID ClientID) error

	// AdoptAgentSessions claims existing sessions for an agent that reconnected.
	// Called when an agent registers - recreates local proxies and updates NodeID.
	// This enables session persistence across pod restarts in distributed deployments.
	AdoptAgentSessions(ctx context.Context, agentID AgentID) error
}