updates

agent/internal/grpc/client.go

@@ -102,3 +102,9 @@ func (c *Client) UploadGeneratedKey(serverID, agentToken, publicKey, label strin
 	}
 	return resp.KeyId, nil
 }
+
+// CommandStream opens a long-lived bidirectional stream for server-pushed commands.
+// The caller controls the stream lifetime via ctx.
+func (c *Client) CommandStream(ctx context.Context) (pb.KeyManager_CommandStreamClient, error) {
+	return c.client.CommandStream(ctx)
+}

agent/internal/grpc/pb/keymanager.pb.go

@@ -42,10 +42,85 @@ type UploadKeyResponse struct {
 	KeyId string `json:"key_id"`
 }
 
+// CommandStream message types
+
+type ServerCommand struct {
+	CommandId   string          `json:"command_id"`
+	GenerateKey *GenerateKeyCmd `json:"generate_key,omitempty"`
+}
+
+type GenerateKeyCmd struct {
+	Label string `json:"label"`
+}
+
+type AgentMessage struct {
+	ServerId   string         `json:"server_id"`
+	AgentToken string         `json:"agent_token"`
+	Ready      *AgentReady    `json:"ready,omitempty"`
+	Result     *CommandResult `json:"result,omitempty"`
+}
+
+type AgentReady struct{}
+
+type CommandResult struct {
+	CommandId string `json:"command_id"`
+	Success   bool   `json:"success"`
+	Message   string `json:"message"`
+}
+
+// CommandStream client-side interface
+
+type KeyManager_CommandStreamClient interface {
+	Send(*AgentMessage) error
+	Recv() (*ServerCommand, error)
+	grpc.ClientStream
+}
+
+type keyManagerCommandStreamClient struct {
+	grpc.ClientStream
+}
+
+func (c *keyManagerCommandStreamClient) Send(m *AgentMessage) error {
+	return c.ClientStream.SendMsg(m)
+}
+
+func (c *keyManagerCommandStreamClient) Recv() (*ServerCommand, error) {
+	m := new(ServerCommand)
+	if err := c.ClientStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
+// CommandStream server-side interface (included for completeness)
+
+type KeyManager_CommandStreamServer interface {
+	Send(*ServerCommand) error
+	Recv() (*AgentMessage, error)
+	grpc.ServerStream
+}
+
+type keyManagerCommandStreamServer struct {
+	grpc.ServerStream
+}
+
+func (s *keyManagerCommandStreamServer) Send(m *ServerCommand) error {
+	return s.ServerStream.SendMsg(m)
+}
+
+func (s *keyManagerCommandStreamServer) Recv() (*AgentMessage, error) {
+	m := new(AgentMessage)
+	if err := s.ServerStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
 type KeyManagerClient interface {
 	Register(ctx context.Context, in *RegisterRequest, opts ...grpc.CallOption) (*RegisterResponse, error)
 	SyncKeys(ctx context.Context, in *SyncRequest, opts ...grpc.CallOption) (*SyncResponse, error)
 	UploadGeneratedKey(ctx context.Context, in *UploadKeyRequest, opts ...grpc.CallOption) (*UploadKeyResponse, error)
+	CommandStream(ctx context.Context, opts ...grpc.CallOption) (KeyManager_CommandStreamClient, error)
 }
 
 type UnimplementedKeyManagerServer struct{}
@@ -91,3 +166,12 @@ func (c *keyManagerClient) UploadGeneratedKey(ctx context.Context, in *UploadKey
 	}
 	return out, nil
 }
+
+func (c *keyManagerClient) CommandStream(ctx context.Context, opts ...grpc.CallOption) (KeyManager_CommandStreamClient, error) {
+	desc := &grpc.StreamDesc{StreamName: "CommandStream", ServerStreams: true, ClientStreams: true}
+	stream, err := c.cc.NewStream(ctx, desc, "/keymanager.v1.KeyManager/CommandStream", opts...)
+	if err != nil {
+		return nil, err
+	}
+	return &keyManagerCommandStreamClient{stream}, nil
+}