From b23805a7b8f4d720ff5ecea2b981a1cc57895283 Mon Sep 17 00:00:00 2001
From: mrhid6 <mrhid6@gmail.com>
Date: Wed, 29 Jul 2026 12:37:37 +0100
Subject: [PATCH] feat: add ProxyStream wire types for agent-relayed console

---
 internal/grpc/client.go        |  4 ++
 internal/grpc/pb/vantage.pb.go | 82 ++++++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/internal/grpc/client.go b/internal/grpc/client.go
index 10171e4..fa6daff 100644
--- a/internal/grpc/client.go
+++ b/internal/grpc/client.go
@@ -151,3 +151,7 @@ func (c *Client) ReportChecks(serverID, agentToken string, results []pb.CheckRes
 func (c *Client) CommandStream(ctx context.Context) (pb.Vantage_CommandStreamClient, error) {
 	return c.client.CommandStream(ctx)
 }
+
+func (c *Client) ProxyStream(ctx context.Context) (pb.Vantage_ProxyStreamClient, error) {
+	return c.client.ProxyStream(ctx)
+}
diff --git a/internal/grpc/pb/vantage.pb.go b/internal/grpc/pb/vantage.pb.go
index 11f9cdb..d96cd29 100644
--- a/internal/grpc/pb/vantage.pb.go
+++ b/internal/grpc/pb/vantage.pb.go
@@ -131,6 +131,32 @@ type ReportChecksResponse struct{}
 
 type ApplyUpdatesCmd struct{}
 
+type OpenProxyCmd struct {
+	ProxyId string `json:"proxy_id"`
+	Port    uint32 `json:"port"`
+}
+
+type ProxyOpen struct {
+	ServerId   string `json:"server_id"`
+	AgentToken string `json:"agent_token"`
+	ProxyId    string `json:"proxy_id"`
+}
+
+type ProxyClose struct {
+	Reason string `json:"reason,omitempty"`
+}
+
+type ProxyClientMsg struct {
+	Open  *ProxyOpen  `json:"open,omitempty"`
+	Data  []byte      `json:"data,omitempty"`
+	Close *ProxyClose `json:"close,omitempty"`
+}
+
+type ProxyServerMsg struct {
+	Data  []byte      `json:"data,omitempty"`
+	Close *ProxyClose `json:"close,omitempty"`
+}
+
 type ServerCommand struct {
 	CommandId        string               `json:"command_id"`
 	GenerateKey      *GenerateKeyCmd      `json:"generate_key,omitempty"`
@@ -139,6 +165,7 @@ type ServerCommand struct {
 	ApplyUpdates     *ApplyUpdatesCmd     `json:"apply_updates,omitempty"`
 	RunStep          *RunStepCmd          `json:"run_step,omitempty"`
 	CleanupWorkspace *CleanupWorkspaceCmd `json:"cleanup_workspace,omitempty"`
+	OpenProxy        *OpenProxyCmd        `json:"open_proxy,omitempty"`
 }
 
 
@@ -254,6 +281,51 @@ func (s *keyManagerCommandStreamServer) Recv() (*AgentMessage, error) {
 	return m, nil
 }
 
+type Vantage_ProxyStreamServer interface {
+	Send(*ProxyServerMsg) error
+	Recv() (*ProxyClientMsg, error)
+	grpc.ServerStream
+}
+
+type vantageProxyStreamServer struct {
+	grpc.ServerStream
+}
+
+func (s *vantageProxyStreamServer) Send(m *ProxyServerMsg) error {
+	return s.ServerStream.SendMsg(m)
+}
+
+func (s *vantageProxyStreamServer) Recv() (*ProxyClientMsg, error) {
+	m := new(ProxyClientMsg)
+	if err := s.ServerStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
+type Vantage_ProxyStreamClient interface {
+	Send(*ProxyClientMsg) error
+	Recv() (*ProxyServerMsg, error)
+	CloseSend() error
+	grpc.ClientStream
+}
+
+type vantageProxyStreamClient struct {
+	grpc.ClientStream
+}
+
+func (c *vantageProxyStreamClient) Send(m *ProxyClientMsg) error {
+	return c.ClientStream.SendMsg(m)
+}
+
+func (c *vantageProxyStreamClient) Recv() (*ProxyServerMsg, error) {
+	m := new(ProxyServerMsg)
+	if err := c.ClientStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
 type VantageClient interface {
 	Register(ctx context.Context, in *RegisterRequest, opts ...grpc.CallOption) (*RegisterResponse, error)
 	SyncKeys(ctx context.Context, in *SyncRequest, opts ...grpc.CallOption) (*SyncResponse, error)
@@ -263,6 +335,7 @@ type VantageClient interface {
 	SyncMonitors(ctx context.Context, in *SyncMonitorsRequest, opts ...grpc.CallOption) (*SyncMonitorsResponse, error)
 	ReportChecks(ctx context.Context, in *ReportChecksRequest, opts ...grpc.CallOption) (*ReportChecksResponse, error)
 	CommandStream(ctx context.Context, opts ...grpc.CallOption) (Vantage_CommandStreamClient, error)
+	ProxyStream(ctx context.Context, opts ...grpc.CallOption) (Vantage_ProxyStreamClient, error)
 }
 
 type UnimplementedVantageServer struct{}
@@ -349,3 +422,12 @@ func (c *keyManagerClient) CommandStream(ctx context.Context, opts ...grpc.CallO
 	}
 	return &vantageCommandStreamClient{stream}, nil
 }
+
+func (c *keyManagerClient) ProxyStream(ctx context.Context, opts ...grpc.CallOption) (Vantage_ProxyStreamClient, error) {
+	desc := &grpc.StreamDesc{StreamName: "ProxyStream", ServerStreams: true, ClientStreams: true}
+	stream, err := c.cc.NewStream(ctx, desc, "/vantage.v1.Vantage/ProxyStream", opts...)
+	if err != nil {
+		return nil, err
+	}
+	return &vantageProxyStreamClient{stream}, nil
+}